fix: codemod gaps (#1246)

* fix: improve codemod logic to cover more edge cases

* fix: handle spyOn and import statements

---------

Co-authored-by: Gergely Békési <gergely.bekesi@ethswarm.org>

Author: slapec93

codemod/transforms/v13.ts

@@ -242,8 +242,259 @@ export function transform(sourceFile: ts.SourceFile, checker: ts.TypeChecker): s
     })
   }
 
+  // ChunkBuilder.hash() now returns a Reference, not a raw Uint8Array.
+  const isChunkBuilderReceiver = (expr: ts.Expression): boolean => {
+    const symbol = checker.getTypeAtLocation(expr).getSymbol()
+
+    if (!symbol || symbol.getName() !== 'ChunkBuilder') {
+      return false
+    }
+
+    return (symbol.getDeclarations() ?? []).some(decl => {
+      if (!ts.isClassDeclaration(decl)) {
+        return false
+      }
+
+      const file = decl.getSourceFile().fileName
+
+      return /[/\\]core-sdk[/\\]/.test(file) || /[/\\]splitter\.(d\.)?ts$/.test(file)
+    })
+  }
+
   const replacements: Array<{ start: number; end: number; text: string }> = []
 
+  // MerkleTree was removed in v13; ChunkSplitter is the replacement.
+  let merkleTreeLocalName: string | null = null
+  let merkleTreeImportNameNode: ts.Identifier | null = null
+
+  for (const stmt of sourceFile.statements) {
+    if (
+      ts.isImportDeclaration(stmt) &&
+      ts.isStringLiteral(stmt.moduleSpecifier) &&
+      stmt.moduleSpecifier.text === '@ethersphere/bee-js' &&
+      stmt.importClause?.namedBindings &&
+      ts.isNamedImports(stmt.importClause.namedBindings)
+    ) {
+      for (const element of stmt.importClause.namedBindings.elements) {
+        const importedNameNode = element.propertyName ?? element.name
+
+        if (importedNameNode.text === 'MerkleTree') {
+          // Aliased imports keep their local name unchanged.
+          replacements.push({
+            start: importedNameNode.getStart(sourceFile),
+            end: importedNameNode.getEnd(),
+            text: 'ChunkSplitter',
+          })
+
+          if (!element.propertyName) {
+            merkleTreeLocalName = element.name.text
+            merkleTreeImportNameNode = element.name
+          }
+        }
+      }
+    }
+
+    // @upcoming/swarm-core is now @ethersphere/core-sdk.
+    if (
+      (ts.isImportDeclaration(stmt) || ts.isExportDeclaration(stmt)) &&
+      stmt.moduleSpecifier &&
+      ts.isStringLiteral(stmt.moduleSpecifier) &&
+      stmt.moduleSpecifier.text === '@upcoming/swarm-core'
+    ) {
+      const quote = stmt.moduleSpecifier.getText(sourceFile)[0]!
+      replacements.push({
+        start: stmt.moduleSpecifier.getStart(sourceFile),
+        end: stmt.moduleSpecifier.getEnd(),
+        text: `${quote}@ethersphere/core-sdk${quote}`,
+      })
+    }
+  }
+
+  // Syntactic fallback: MerkleTree no longer type-checks, so isChunkBuilderReceiver alone
+  // won't find it.
+  const unwrap = (expr: ts.Expression): ts.Expression => {
+    while (ts.isParenthesizedExpression(expr) || ts.isAwaitExpression(expr)) {
+      expr = expr.expression
+    }
+
+    return expr
+  }
+
+  const isChunkSplitterRootCall = (expr: ts.Expression): boolean => {
+    const inner = unwrap(expr)
+
+    return (
+      merkleTreeLocalName !== null &&
+      ts.isCallExpression(inner) &&
+      ts.isPropertyAccessExpression(inner.expression) &&
+      ts.isIdentifier(inner.expression.expression) &&
+      inner.expression.expression.text === merkleTreeLocalName &&
+      (inner.expression.name.text === 'root' || inner.expression.name.text === 'finalize')
+    )
+  }
+
+  const chunkBuilderVarNames = new Set<string>()
+
+  const collectChunkBuilderVars = (node: ts.Node): void => {
+    if (
+      ts.isVariableDeclaration(node) &&
+      ts.isIdentifier(node.name) &&
+      node.initializer &&
+      isChunkSplitterRootCall(node.initializer)
+    ) {
+      chunkBuilderVarNames.add(node.name.text)
+    }
+
+    ts.forEachChild(node, collectChunkBuilderVars)
+  }
+
+  if (merkleTreeLocalName) {
+    collectChunkBuilderVars(sourceFile)
+  }
+
+  // bee-js's MantarayNode.collect()/.find()/.findClosest() return core-sdk's MantarayNode,
+  // not its own.
+  const isMantarayNodeType = (type: ts.Type, fileTest: RegExp): boolean => {
+    const symbol = type.getSymbol()
+
+    if (!symbol || symbol.getName() !== 'MantarayNode') {
+      return false
+    }
+
+    return (symbol.getDeclarations() ?? []).some(
+      decl => ts.isClassDeclaration(decl) && fileTest.test(decl.getSourceFile().fileName),
+    )
+  }
+  const isBeeMantarayNodeType = (type: ts.Type): boolean => isMantarayNodeType(type, /[/\\]bee-js[/\\]/)
+  const isCoreMantarayNodeType = (type: ts.Type): boolean => isMantarayNodeType(type, /[/\\]core-sdk[/\\]/)
+
+  // Tracks for-of loop vars whose element type is core-sdk's MantarayNode.
+  const coreFlavoredNames = new Set<string>()
+
+  const collectCoreFlavoredBindings = (node: ts.Node): void => {
+    if (ts.isForOfStatement(node) && ts.isVariableDeclarationList(node.initializer)) {
+      const decl = node.initializer.declarations[0]
+
+      if (decl && ts.isIdentifier(decl.name) && isCoreMantarayNodeType(checker.getTypeAtLocation(decl.name))) {
+        coreFlavoredNames.add(decl.name.text)
+      }
+    }
+
+    ts.forEachChild(node, collectCoreFlavoredBindings)
+  }
+  collectCoreFlavoredBindings(sourceFile)
+
+  let needsCoreMantarayNodeImport = false
+
+  // Rewrites a parameter's type when called with a core-flavored argument.
+  const rewriteMantarayNodeParameters = (node: ts.Node): void => {
+    if (ts.isCallExpression(node)) {
+      node.arguments.forEach((arg, index) => {
+        if (!ts.isIdentifier(arg) || !coreFlavoredNames.has(arg.text)) {
+          return
+        }
+
+        const signature = checker.getResolvedSignature(node)
+        const decl = signature?.getDeclaration()
+        const param = decl && 'parameters' in decl ? decl.parameters[index] : undefined
+
+        if (
+          param?.type &&
+          ts.isTypeReferenceNode(param.type) &&
+          isBeeMantarayNodeType(checker.getTypeFromTypeNode(param.type))
+        ) {
+          replacements.push({
+            start: param.type.getStart(sourceFile),
+            end: param.type.getEnd(),
+            text: 'CoreMantarayNode',
+          })
+          needsCoreMantarayNodeImport = true
+        }
+      })
+    }
+
+    ts.forEachChild(node, rewriteMantarayNodeParameters)
+  }
+
+  // Rewrites a Map/Array/Set's value type when a core-flavored value is inserted.
+  const rewriteMantarayNodeContainers = (node: ts.Node): void => {
+    if (
+      ts.isCallExpression(node) &&
+      ts.isPropertyAccessExpression(node.expression) &&
+      ['set', 'push', 'add'].includes(node.expression.name.text)
+    ) {
+      const valueArg = node.expression.name.text === 'set' ? node.arguments[1] : node.arguments[0]
+
+      if (valueArg && ts.isIdentifier(valueArg) && coreFlavoredNames.has(valueArg.text)) {
+        const receiverSymbol = checker.getSymbolAtLocation(node.expression.expression)
+        const receiverDecl = receiverSymbol?.declarations?.find(ts.isVariableDeclaration)
+
+        // Type args may be on the declaration or the constructor call.
+        const typeArguments =
+          receiverDecl?.type && ts.isTypeReferenceNode(receiverDecl.type)
+            ? receiverDecl.type.typeArguments
+            : receiverDecl?.initializer && ts.isNewExpression(receiverDecl.initializer)
+              ? receiverDecl.initializer.typeArguments
+              : undefined
+
+        if (typeArguments) {
+          for (const typeArg of typeArguments) {
+            if (ts.isTypeReferenceNode(typeArg) && isBeeMantarayNodeType(checker.getTypeFromTypeNode(typeArg))) {
+              replacements.push({
+                start: typeArg.getStart(sourceFile),
+                end: typeArg.getEnd(),
+                text: 'CoreMantarayNode',
+              })
+              needsCoreMantarayNodeImport = true
+            }
+          }
+        }
+      }
+    }
+
+    ts.forEachChild(node, rewriteMantarayNodeContainers)
+  }
+
+  if (coreFlavoredNames.size > 0) {
+    rewriteMantarayNodeParameters(sourceFile)
+    rewriteMantarayNodeContainers(sourceFile)
+  }
+
+  if (needsCoreMantarayNodeImport) {
+    const coreSdkImport = sourceFile.statements.find(
+      (stmt): stmt is ts.ImportDeclaration =>
+        ts.isImportDeclaration(stmt) &&
+        ts.isStringLiteral(stmt.moduleSpecifier) &&
+        stmt.moduleSpecifier.text === '@ethersphere/core-sdk',
+    )
+
+    if (coreSdkImport?.importClause?.namedBindings && ts.isNamedImports(coreSdkImport.importClause.namedBindings)) {
+      const elements = coreSdkImport.importClause.namedBindings.elements
+      const lastSpecifier = elements[elements.length - 1]
+
+      if (lastSpecifier) {
+        replacements.push({
+          start: lastSpecifier.getEnd(),
+          end: lastSpecifier.getEnd(),
+          text: ', MantarayNode as CoreMantarayNode',
+        })
+      }
+    } else {
+      const beeJsImport = sourceFile.statements.find(
+        (stmt): stmt is ts.ImportDeclaration =>
+          ts.isImportDeclaration(stmt) &&
+          ts.isStringLiteral(stmt.moduleSpecifier) &&
+          stmt.moduleSpecifier.text === '@ethersphere/bee-js',
+      )
+      const insertAt = beeJsImport ? beeJsImport.getEnd() : 0
+      replacements.push({
+        start: insertAt,
+        end: insertAt,
+        text: `\nimport { MantarayNode as CoreMantarayNode } from '@ethersphere/core-sdk'`,
+      })
+    }
+  }
+
   const visit = (node: ts.Node): void => {
     // Any `<bee>.<method>` access — called or not — so bare method references migrate too.
     if (ts.isPropertyAccessExpression(node) && ts.isIdentifier(node.name)) {
@@ -261,6 +512,66 @@ export function transform(sourceFile: ts.SourceFile, checker: ts.TypeChecker): s
       }
     }
 
+    // jest.spyOn uses a string literal, so the rewrite above doesn't catch it.
+    if (
+      ts.isCallExpression(node) &&
+      ts.isPropertyAccessExpression(node.expression) &&
+      node.expression.name.text === 'spyOn' &&
+      node.arguments.length >= 2
+    ) {
+      const [receiverArg, methodArg] = node.arguments
+
+      if (receiverArg && methodArg && ts.isStringLiteralLike(methodArg)) {
+        const mapping = METHOD_MAP[methodArg.text]
+
+        if (mapping && isBeeReceiver(receiverArg)) {
+          const quote = methodArg.getText(sourceFile)[0]!
+          replacements.push({
+            start: receiverArg.getStart(sourceFile),
+            end: receiverArg.getEnd(),
+            text: `${receiverArg.getText(sourceFile)}.${mapping.namespace}`,
+          })
+          replacements.push({
+            start: methodArg.getStart(sourceFile),
+            end: methodArg.getEnd(),
+            text: `${quote}${mapping.newName}${quote}`,
+          })
+        }
+      }
+    }
+
+    // ChunkBuilder.hash() now returns a Reference; append the conversion back.
+    if (
+      merkleTreeLocalName &&
+      ts.isCallExpression(node) &&
+      ts.isPropertyAccessExpression(node.expression) &&
+      node.expression.name.text === 'hash' &&
+      node.arguments.length === 0 &&
+      (isChunkSplitterRootCall(node.expression.expression) ||
+        (ts.isIdentifier(node.expression.expression) && chunkBuilderVarNames.has(node.expression.expression.text)) ||
+        isChunkBuilderReceiver(node.expression.expression))
+    ) {
+      replacements.push({
+        start: node.getEnd(),
+        end: node.getEnd(),
+        text: '.toUint8Array()',
+      })
+    }
+
+    // Renames remaining MerkleTree usages (import itself already handled above).
+    if (
+      merkleTreeLocalName &&
+      ts.isIdentifier(node) &&
+      node.text === merkleTreeLocalName &&
+      node !== merkleTreeImportNameNode
+    ) {
+      replacements.push({
+        start: node.getStart(sourceFile),
+        end: node.getEnd(),
+        text: 'ChunkSplitter',
+      })
+    }
+
     ts.forEachChild(node, visit)
   }
   visit(sourceFile)