diff --git a/src/index.js b/src/index.js
index 086b983..ccbcc1e 100644
--- a/src/index.js
+++ b/src/index.js
@@ -315,6 +315,8 @@ export default function ({ Plugin, types: t }) {
             return;
           }
 
+          this.state[foundJSXKey] = false;
+
           // No-op for `render` methods in this visitor
           if (parent.key && parent.key.name === 'render') {
             return;
@@ -342,6 +344,10 @@ export default function ({ Plugin, types: t }) {
             );
           }
 
+          if (t.isFunctionDeclaration(node)) {
+            // console.log(node.body.body.filter(t.isReturnStatement));
+          }
+
           return [
             node,
             t.assignmentExpression('=', bindingId,
diff --git a/test/fixtures/pure-function/expected.js b/test/fixtures/pure-function/expected.js
index 2566f93..64425f3 100644
--- a/test/fixtures/pure-function/expected.js
+++ b/test/fixtures/pure-function/expected.js
@@ -78,6 +78,7 @@ exports.NamedExportFunction = NamedExportFunction = _wrapComponent('_$NamedExpor
 function DefaultExportFunction() {
   return _react2['default'].createElement('div', null);
 }
+
 DefaultExportFunction = _wrapComponent('_$DefaultExportFunction')(DefaultExportFunction)
 var ModulePattern = {
   NoProps: _wrapComponent('_$NoProps')(function () {