refactor: optimize and cache variable lookup

src/StateSnapshot.js

@@ -285,32 +285,50 @@ function isInAssignmentExpression(node) {
   )
 }
 
+const whichCache = new WeakMap()
+
 function which(name, scope) {
-  let kind = null
+  if (!scope) return null
 
-  if (!scope) return kind
+  let cachedForScope = whichCache.get(scope)
+  if (!cachedForScope) {
+    cachedForScope = new Map()
+    whichCache.set(scope, cachedForScope)
+  }
+  if (cachedForScope.has(name)) return cachedForScope.get(name)
+
+  let kind = null
+  for (const variable of scope.variables) {
+    if (variable.name !== name) continue
 
-  scope.variables.forEach((variable) => {
     const def = variable.defs[0]
-    if (!def || variable.name !== name) return
+    if (!def) continue
 
     const init = def.node.init
-    if (!init) return
+    if (!init) continue
+
     if (init.type === 'Identifier') {
-      return (kind = which(init.name, scope))
-    } else if (init.type === 'CallExpression' && init.callee.name === 'proxy') {
-      return (kind = 'state')
-    } else if (
-      init.type === 'CallExpression' &&
-      init.callee.name === 'useSnapshot'
-    ) {
-      return (kind = 'snapshot')
+      kind = which(init.name, scope)
+    } else if (init.type === 'CallExpression') {
+      const { name: calleeName } = init.callee
+      if (calleeName === 'proxy') {
+        kind = 'state'
+      } else if (calleeName === 'useSnapshot') {
+        kind = 'snapshot'
+      }
     }
-  })
-  if (!kind && scope.upper) return (kind = which(name, scope.upper))
 
+    if (kind) break
+  }
+
+  if (!kind && scope.upper) {
+    kind = which(name, scope.upper)
+  }
+
+  cachedForScope.set(name, kind)
   return kind
 }
+
 function isSameMemmberExpression(first, second) {
   if (!first || !second) return false
   if (
@@ -335,6 +353,7 @@ function isSameMemmberExpression(first, second) {
   }
   return false
 }
+
 function isUsedInUseProxy(node, scope) {
   let isUsed = false
   if (!scope) return isUsed

src/lib/utils.js

@@ -67,6 +67,7 @@ export function nearestCalleeName(node) {
   return nearestCalleeName(node.parent)
 }
 
+const parentOfNodeTypeCache = new WeakMap()
 /**
  * @param {any} node ASTNode to start from
  * @param {string} nodeType the type of ASTNode to look for
@@ -88,12 +89,25 @@ export function getParentOfNodeType(node, nodeType) {
     return null
   }
 
-  if (node.parent && node.parent.type !== nodeType) {
-    return getParentOfNodeType(node.parent, nodeType)
-  } else if (node.parent && node.parent.type === nodeType) {
-    return node.parent
+  let cacheForNode = parentOfNodeTypeCache.get(node)
+  if (cacheForNode && cacheForNode[nodeType] !== undefined) {
+    return cacheForNode[nodeType]
   }
-  return null
+
+  let result = null
+  if (node.parent.type !== nodeType) {
+    result = getParentOfNodeType(node.parent, nodeType)
+  } else if (node.parent.type === nodeType) {
+    result = node.parent
+  }
+
+  if (!cacheForNode) {
+    cacheForNode = {}
+    parentOfNodeTypeCache.set(node, cacheForNode)
+  }
+  cacheForNode[nodeType] = result
+
+  return result
 }
 
 /**