From 563d5d9c344031d6d6685358391c13dba258d3ce Mon Sep 17 00:00:00 2001
From: proteriax <8125011+proteriax@users.noreply.github.com>
Date: Thu, 22 Jul 2021 09:54:10 -0400
Subject: [PATCH] Prevent infinite visiting

---
 src/index.ts | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index ef62397..fc89d9c 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -2,9 +2,10 @@ import type * as babel from "@babel/core"
 import type * as t from "@babel/types"
 
 interface State {
+  visited: WeakSet<t.ClassDeclaration>
   getTSLib(): t.ImportSpecifier[]
-  getParam(): t.Identifier
-  getDecorate(): t.Identifier
+  __param(): t.Identifier
+  __decorate(): t.Identifier
 }
 
 export default ({ types: t }: typeof babel): babel.PluginObj<State> => ({
@@ -34,12 +35,16 @@ export default ({ types: t }: typeof babel): babel.PluginObj<State> => ({
         const getParam = createGet("__param")
         const getDecorate = createGet("__decorate")
 
-        state.getParam = () => t.cloneNode(getParam())
-        state.getDecorate = () => t.cloneNode(getDecorate())
+        state.__param = () => t.cloneNode(getParam())
+        state.__decorate = () => t.cloneNode(getDecorate())
+        state.visited = new WeakSet()
       },
     },
 
     ClassDeclaration(path, state) {
+      if (state.visited.has(path.node)) return
+      state.visited.add(path.node)
+
       const modifiers: (t.Expression | t.Statement)[] = []
       const classDecorators: t.Expression[] = []
       const prefixes: t.Statement[] = []
@@ -59,7 +64,7 @@ export default ({ types: t }: typeof babel): babel.PluginObj<State> => ({
           case "ClassProperty":
             if (!child.decorators) continue
             modifiers.push(
-              t.callExpression(state.getDecorate(), [
+              t.callExpression(state.__decorate(), [
                 t.arrayExpression(child.decorators.map(d => d.expression)),
                 t.memberExpression(classID, t.identifier("prototype")),
                 t.isIdentifier(child.key) ? t.stringLiteral(child.key.name) : child.key,
@@ -83,9 +88,7 @@ export default ({ types: t }: typeof babel): babel.PluginObj<State> => ({
                 list.push(
                   ...param.decorators
                     .map(d => d.expression)
-                    .map(e =>
-                      t.callExpression(state.getParam(), [t.numericLiteral(i), e])
-                    )
+                    .map(e => t.callExpression(state.__param(), [t.numericLiteral(i), e]))
                 )
                 param.decorators = null
               })
@@ -108,7 +111,7 @@ export default ({ types: t }: typeof babel): babel.PluginObj<State> => ({
                 }
 
                 modifiers.push(
-                  t.callExpression(state.getDecorate(), [
+                  t.callExpression(state.__decorate(), [
                     t.arrayExpression(list),
                     t.memberExpression(classID, t.identifier("prototype")),
                     key,
@@ -134,7 +137,7 @@ export default ({ types: t }: typeof babel): babel.PluginObj<State> => ({
           t.assignmentExpression(
             "=",
             classID,
-            t.callExpression(state.getDecorate(), [
+            t.callExpression(state.__decorate(), [
               t.arrayExpression(classDecorators),
               classID,
             ])