feat: add type narrow to support better type check

Author: HerrCai0907

src/compiler.ts

@@ -88,7 +88,8 @@ import {
   PropertyPrototype,
   IndexSignature,
   File,
-  mangleInternalName
+  mangleInternalName,
+  TypedElement
 } from "./program";
 
 import {
@@ -210,6 +211,7 @@ import {
 import {
   ShadowStackPass
 } from "./passes/shadowstack";
+import { TypeNarrowInfo } from "./conditionInfo";
 
 /** Compiler options. */
 export class Options {
@@ -2665,12 +2667,17 @@ export class Compiler extends DiagnosticEmitter {
     // )                 └┬─────────┴─┘
     // ...              ┌◄┘
 
+    this.currentFlow.startCompileCondition();
+
     // Precompute the condition (always executes)
     var condExpr = this.makeIsTrueish(
       this.compileExpression(statement.condition, Type.bool),
       this.currentType,
       statement.condition
     );
+
+    const conditionInfoContainer = this.currentFlow.stopCompileCondition();
+
     var condKind = this.evaluateCondition(condExpr);
 
     // Shortcut if the condition is constant
@@ -2700,11 +2707,13 @@ export class Compiler extends DiagnosticEmitter {
     var thenFlow = flow.fork();
     this.currentFlow = thenFlow;
     thenFlow.inheritNonnullIfTrue(condExpr);
+    conditionInfoContainer.trueInfo.forEach((info) => info.apply());
     if (ifTrue.kind == NodeKind.BLOCK) {
       this.compileStatements((<BlockStatement>ifTrue).statements, false, thenStmts);
     } else {
       thenStmts.push(this.compileStatement(ifTrue));
     }
+    conditionInfoContainer.trueInfo.forEach((info) => info.recover());
     var thenTerminates = thenFlow.isAny(FlowFlags.TERMINATES | FlowFlags.BREAKS);
     if (thenTerminates) {
       thenStmts.push(module.unreachable());
@@ -2718,11 +2727,13 @@ export class Compiler extends DiagnosticEmitter {
       let elseFlow = flow.fork();
       this.currentFlow = elseFlow;
       elseFlow.inheritNonnullIfFalse(condExpr);
+      conditionInfoContainer.falseInfo.forEach((info) => info.apply());
       if (ifFalse.kind == NodeKind.BLOCK) {
         this.compileStatements((<BlockStatement>ifFalse).statements, false, elseStmts);
       } else {
         elseStmts.push(this.compileStatement(ifFalse));
       }
+      conditionInfoContainer.falseInfo.forEach((info) => info.recover());
       let elseTerminates = elseFlow.isAny(FlowFlags.TERMINATES | FlowFlags.BREAKS);
       if (elseTerminates) {
         elseStmts.push(module.unreachable());
@@ -5846,6 +5857,7 @@ export class Compiler extends DiagnosticEmitter {
     var flow = this.currentFlow;
     var target = resolver.lookupExpression(expression, flow); // reports
     if (!target) return this.module.unreachable();
+    if (target instanceof TypedElement) (<TypedElement>target).recoverType();
     var thisExpression = resolver.currentThisExpression;
     var elementExpression = resolver.currentElementExpression;
 
@@ -7903,6 +7915,14 @@ export class Compiler extends DiagnosticEmitter {
       this.currentType = Type.bool;
       return this.module.unreachable();
     }
+    if (expression.expression.kind == NodeKind.IDENTIFIER) {
+      const identifier = <IdentifierExpression>expression.expression;
+      const element = flow.lookup(identifier.text);
+      if (element && element instanceof TypedElement) {
+        const typedElement = <TypedElement>element;
+        flow.addConditionInfo(new TypeNarrowInfo(typedElement, expectedType));
+      }
+    }
     return this.makeInstanceofType(expression, expectedType);
   }
 

src/conditionInfo.ts

@@ -0,0 +1,31 @@
+import { TypedElement } from "./program";
+import { Type } from "./types";
+
+export class ConditionInfoContainer {
+  trueInfo: ConditionInfo[] = [];
+  falseInfo: ConditionInfo[] = [];
+
+  switchTrueAndFalse(): void {
+    let tmp = this.trueInfo;
+    this.trueInfo = this.falseInfo
+    this.falseInfo = tmp
+  }
+}
+
+export class ConditionInfo {
+  apply(): void {}
+  recover(): void {}
+}
+
+export class TypeNarrowInfo extends ConditionInfo {
+  constructor(public typedElement: TypedElement, public narrowedType: Type) {
+    super();
+  }
+
+  override apply(): void {
+    this.typedElement.narrowType(this.narrowedType);
+  }
+  override recover(): void {
+    this.typedElement.recoverType();
+  }
+}

src/flow.ts

@@ -90,6 +90,7 @@ import {
 import {
   BuiltinNames
 } from "./builtins";
+import { ConditionInfo, ConditionInfoContainer } from "./conditionInfo";
 
 /** Control flow flags indicating specific conditions. */
 export const enum FlowFlags {
@@ -1221,6 +1222,25 @@ export class Flow {
     }
   }
 
+  conditionInfoStack: ConditionInfoContainer[] = []
+  startCompileCondition(): void {
+    this.conditionInfoStack.push(new ConditionInfoContainer());
+  }
+  stopCompileCondition(): ConditionInfoContainer {
+    return assert(this.conditionInfoStack.pop());
+  }
+  addConditionInfo(conditionInfo: ConditionInfo, branch: bool = true): void {
+    if (this.conditionInfoStack.length > 0) {
+      const infoContainer =
+        this.conditionInfoStack[this.conditionInfoStack.length - 1];
+      if (branch) {
+        infoContainer.trueInfo.push(conditionInfo);
+      } else {
+        infoContainer.falseInfo.push(conditionInfo);
+      }
+    }    
+  }
+
   /**
    * Tests if an expression can possibly overflow in the context of this flow. Assumes that the
    * expression might already have overflown and returns `false` only if the operation neglects

src/program.ts

@@ -2992,6 +2992,7 @@ export abstract class TypedElement extends DeclaredElement {
 
   /** Resolved type. Set once `is(RESOLVED)`, otherwise void. */
   type: Type = Type.void;
+  typeStack: Type[] = [];
 
   constructor(
     /** Specific element kind. */
@@ -3017,6 +3018,15 @@ export abstract class TypedElement extends DeclaredElement {
     this.type = type;
     this.set(CommonFlags.RESOLVED);
   }
+  narrowType(type: Type): void {
+    if (type.isClass) {
+      this.typeStack.push(this.type);
+      this.type = type;
+    }
+  }
+  recoverType(): void {
+    this.type = this.typeStack.pop() || this.type;
+  }
 }
 
 /** A file representing the implicit top-level namespace of a source. */

tests/compiler/typenarrow-error.json

@@ -0,0 +1,8 @@
+{
+  "asc_flags": [
+  ],
+  "stderr": [
+    "TS2339: Property 'foo' does not exist on type 'typenarrow-error/A",
+    "EOF"
+  ]
+}

tests/compiler/typenarrow-error.ts

@@ -0,0 +1,13 @@
+class A {}
+class B extends A {
+  foo(): void {}
+}
+
+let t = new A();
+if (t instanceof B) {
+  t.foo();
+  t = new A();
+  t.foo();
+}
+
+ERROR("EOF");