From 29a3f2275622ece985b07db593031b1c1eeed940 Mon Sep 17 00:00:00 2001
From: Alexis Munsayac <alexis.munsayac@gmail.com>
Date: Fri, 17 Sep 2021 21:30:50 +0800
Subject: [PATCH] Fix regression for `AssignmentExpression` support of `memo`

---
 .../src/index.ts                              | 29 +++++++++++--------
 .../test/index.js                             | 11 +------
 2 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/packages/babel-plugin-solid-reactivity/src/index.ts b/packages/babel-plugin-solid-reactivity/src/index.ts
index c81afc7..9270b00 100644
--- a/packages/babel-plugin-solid-reactivity/src/index.ts
+++ b/packages/babel-plugin-solid-reactivity/src/index.ts
@@ -300,6 +300,7 @@ function memoSingleExpression(
   path: NodePath<t.LabeledStatement>,
   memoIdentifier: t.Identifier,
   stateIdentifier: t.Expression,
+  replace = false,
 ): void {
   const readIdentifier = path.scope.generateUidIdentifier(memoIdentifier.name);
   const expr = t.variableDeclaration(
@@ -324,7 +325,11 @@ function memoSingleExpression(
     )],
   );
 
-  path.insertAfter(expr);
+  if (replace) {
+    path.replaceWith(expr);
+  } else {
+    path.insertAfter(expr);
+  }
 
   const parent = path.scope.path;
   if (parent) {
@@ -435,19 +440,19 @@ function memoExpression(
   body: t.Statement,
 ): void {
   if (t.isExpressionStatement(body)) {
-    if (!t.isAssignmentExpression(body.expression)) {
+    if (t.isAssignmentExpression(body.expression)) {
+      if (body.expression.operator !== '=') {
+        throw new Error('Invalid assignment expression operator');
+      }
+      const leftExpr = body.expression.left;
+      const rightExpr = body.expression.right;
+      if (!t.isIdentifier(leftExpr)) {
+        throw new Error('Expected identifier');
+      }
+      memoSingleExpression(hooks, path, leftExpr, rightExpr, true);
+    } else {
       throw new Error('Expected assignment expression');
     }
-    if (body.expression.operator !== '=') {
-      throw new Error('Invalid assignment expression operator');
-    }
-    const leftExpr = body.expression.left;
-    const rightExpr = body.expression.right;
-    if (!t.isIdentifier(leftExpr)) {
-      throw new Error('Expected identifier');
-    }
-    memoSingleExpression(hooks, path, leftExpr, rightExpr);
-    path.remove();
   } else if (t.isVariableDeclaration(body)) {
     for (let i = body.declarations.length - 1; i >= 0; i -= 1) {
       const declarator = body.declarations[i];
diff --git a/packages/babel-plugin-solid-reactivity/test/index.js b/packages/babel-plugin-solid-reactivity/test/index.js
index fdf62cf..f92debe 100644
--- a/packages/babel-plugin-solid-reactivity/test/index.js
+++ b/packages/babel-plugin-solid-reactivity/test/index.js
@@ -2,16 +2,7 @@ const babel = require('@babel/core');
 const plugin = require('../dist/cjs');
 
 const code = `
-function Example() {
-  signal: var x = 0, y = 1;
-
-  while (true) {
-    let y = x;
-  }
-
-  x += y = 2;
-  x++;
-}
+memo: var x = 0, y = 0;
 `;
 babel.transformAsync(code, {
   plugins: [