fix(migrations): inject migration not inserting generated code after super call in some cases (#58393)

crisbeto · alxhub · commit 7a65cdd911cb · 2024-10-28T12:39:03.000-07:00
Fixes an issue where the `inject` migration was generating and attempting to insert code after a `super` call, but the string buffering implementation was dropping it if the statement right after the `super` call was deleted as a result of the migration. PR Close #58393
diff --git a/packages/core/schematics/ng-generate/inject-migration/migration.ts b/packages/core/schematics/ng-generate/inject-migration/migration.ts
@@ -224,7 +224,16 @@ function migrateClass(
   }
 
   if (afterSuper.length > 0 && superCall !== null) {
-    tracker.insertText(sourceFile, superCall.getEnd() + 1, `\n${afterSuper.join('\n')}\n`);
+    // Note that if we can, we should insert before the next statement after the `super` call,
+    // rather than after the end of it. Otherwise the string buffering implementation may drop
+    // the text if the statement after the `super` call is being deleted. This appears to be because
+    // the full start of the next statement appears to always be the end of the `super` call plus 1.
+    const nextStatement = getNextPreservedStatement(superCall, removedStatements);
+    tracker.insertText(
+      sourceFile,
+      nextStatement ? nextStatement.getFullStart() : superCall.getEnd() + 1,
+      `\n${afterSuper.join('\n')}\n`,
+    );
   }
 
   // Need to resolve this once all constructor signatures have been removed.
@@ -676,6 +685,36 @@ function canRemoveConstructor(
   );
 }
 
+/**
+ * Gets the next statement after a node that *won't* be deleted by the migration.
+ * @param startNode Node from which to start the search.
+ * @param removedStatements Statements that have been removed by the migration.
+ * @returns
+ */
+function getNextPreservedStatement(
+  startNode: ts.Node,
+  removedStatements: Set<ts.Statement>,
+): ts.Statement | null {
+  const body = closestNode(startNode, ts.isBlock);
+  const closestStatement = closestNode(startNode, ts.isStatement);
+  if (body === null || closestStatement === null) {
+    return null;
+  }
+
+  const index = body.statements.indexOf(closestStatement);
+  if (index === -1) {
+    return null;
+  }
+
+  for (let i = index + 1; i < body.statements.length; i++) {
+    if (!removedStatements.has(body.statements[i])) {
+      return body.statements[i];
+    }
+  }
+
+  return null;
+}
+
 /**
  * Applies the internal-specific migrations to a class.
  * @param node Class being migrated.
diff --git a/packages/core/schematics/test/inject_migration_spec.ts b/packages/core/schematics/test/inject_migration_spec.ts
@@ -1941,5 +1941,50 @@ describe('inject migration', () => {
         `}`,
       ]);
     });
+
+    it('should be able to insert statements after the `super` call when running in internal migration mode', async () => {
+      writeFile(
+        '/dir.ts',
+        [
+          `import { Directive, Inject, ElementRef } from '@angular/core';`,
+          `import { Foo } from 'foo';`,
+          `import { Parent } from './parent';`,
+          ``,
+          `@Directive()`,
+          `class MyDir extends Parent {`,
+          `  private value: number;`,
+          ``,
+          `  constructor(private foo: Foo, readonly elementRef: ElementRef) {`,
+          `    super();`,
+          `    this.value = this.foo.getValue();`,
+          `    console.log(elementRef.nativeElement.tagName);`,
+          `  }`,
+          `}`,
+        ].join('\n'),
+      );
+
+      await runInternalMigration();
+
+      expect(tree.readContent('/dir.ts').split('\n')).toEqual([
+        `import { Directive, ElementRef, inject } from '@angular/core';`,
+        `import { Foo } from 'foo';`,
+        `import { Parent } from './parent';`,
+        ``,
+        `@Directive()`,
+        `class MyDir extends Parent {`,
+        `  private foo = inject(Foo);`,
+        `  readonly elementRef = inject(ElementRef);`,
+        ``,
+        `  private value: number = this.foo.getValue();`,
+        ``,
+        `  constructor() {`,
+        `    super();`,
+        `    const elementRef = this.elementRef;`,
+        ``,
+        `    console.log(elementRef.nativeElement.tagName);`,
+        `  }`,
+        `}`,
+      ]);
+    });
   });
 });
