Implement non-null assertions (#443)

src/ast.ts

@@ -116,10 +116,11 @@ export function nodeIsConstantValue(kind: NodeKind): bool {
 export function nodeIsCallable(kind: NodeKind): bool {
   switch (kind) {
     case NodeKind.IDENTIFIER:
+    case NodeKind.ASSERTION: // if kind=NONNULL
     case NodeKind.CALL:
     case NodeKind.ELEMENTACCESS:
-    case NodeKind.PROPERTYACCESS:
-    case NodeKind.PARENTHESIZED: return true;
+    case NodeKind.PARENTHESIZED:
+    case NodeKind.PROPERTYACCESS: return true;
   }
   return false;
 }
@@ -286,14 +287,14 @@ export abstract class Node {
   static createAssertionExpression(
     assertionKind: AssertionKind,
     expression: Expression,
-    toType: CommonTypeNode,
+    toType: CommonTypeNode | null,
     range: Range
   ): AssertionExpression {
     var expr = new AssertionExpression();
     expr.range = range;
     expr.assertionKind = assertionKind;
     expr.expression = expression; expression.parent = expr;
-    expr.toType = toType; toType.parent = expr;
+    expr.toType = toType; if (toType) toType.parent = expr;
     return expr;
   }
 
@@ -1282,7 +1283,8 @@ export class ArrayLiteralExpression extends LiteralExpression {
 /** Indicates the kind of an assertion. */
 export enum AssertionKind {
   PREFIX,
-  AS
+  AS,
+  NONNULL
 }
 
 /** Represents an assertion expression. */
@@ -1294,7 +1296,7 @@ export class AssertionExpression extends Expression {
   /** Expression being asserted. */
   expression: Expression;
   /** Target type. */
-  toType: CommonTypeNode;
+  toType: CommonTypeNode | null;
 }
 
 /** Represents a binary expression. */

src/compiler.ts

@@ -91,6 +91,7 @@ import {
   Source,
   Range,
   DecoratorKind,
+  AssertionKind,
 
   Statement,
   BlockStatement,
@@ -2704,12 +2705,25 @@ export class Compiler extends DiagnosticEmitter {
   }
 
   compileAssertionExpression(expression: AssertionExpression, contextualType: Type): ExpressionRef {
-    var toType = this.resolver.resolveType( // reports
-      expression.toType,
-      this.currentFunction.flow.contextualTypeArguments
-    );
-    if (!toType) return this.module.createUnreachable();
-    return this.compileExpression(expression.expression, toType, ConversionKind.EXPLICIT, WrapMode.NONE);
+    switch (expression.assertionKind) {
+      case AssertionKind.PREFIX:
+      case AssertionKind.AS: {
+        let toType = this.resolver.resolveType( // reports
+          assert(expression.toType),
+          this.currentFunction.flow.contextualTypeArguments
+        );
+        if (!toType) return this.module.createUnreachable();
+        return this.compileExpression(expression.expression, toType, ConversionKind.EXPLICIT, WrapMode.NONE);
+      }
+      case AssertionKind.NONNULL: {
+        assert(!expression.toType);
+        let expr = this.compileExpressionRetainType(expression.expression, contextualType, WrapMode.NONE);
+        this.currentType = this.currentType.nonNullableType;
+        return expr;
+      }
+      default: assert(false);
+    }
+    return this.module.createUnreachable();
   }
 
   private f32ModInstance: Function | null = null;

src/extra/ast.ts

@@ -477,15 +477,26 @@ export class ASTBuilder {
 
   visitAssertionExpression(node: AssertionExpression): void {
     var sb = this.sb;
-    if (node.assertionKind == AssertionKind.PREFIX) {
-      sb.push("<");
-      this.visitTypeNode(node.toType);
-      sb.push(">");
-      this.visitNode(node.expression);
-    } else {
-      this.visitNode(node.expression);
-      sb.push(" as ");
-      this.visitTypeNode(node.toType);
+    switch (node.assertionKind) {
+      case AssertionKind.PREFIX: {
+        sb.push("<");
+        this.visitTypeNode(assert(node.toType));
+        sb.push(">");
+        this.visitNode(node.expression);
+        break;
+      }
+      case AssertionKind.AS: {
+        this.visitNode(node.expression);
+        sb.push(" as ");
+        this.visitTypeNode(assert(node.toType));
+        break;
+      }
+      case AssertionKind.NONNULL: {
+        this.visitNode(node.expression);
+        sb.push("!");
+        break;
+      }
+      default: assert(false);
     }
   }
 

src/parser.ts

@@ -3472,6 +3472,15 @@ export class Parser extends DiagnosticEmitter {
           );
           break;
         }
+        case Token.EXCLAMATION: {
+          expr = Node.createAssertionExpression(
+            AssertionKind.NONNULL,
+            expr,
+            null,
+            tn.range(startPos, tn.pos)
+          );
+          break;
+        }
         // InstanceOfExpression
         case Token.INSTANCEOF: {
           let isType = this.parseType(tn); // reports
@@ -3835,7 +3844,8 @@ function determinePrecedence(kind: Token): Precedence {
     case Token.MINUS_MINUS: return Precedence.UNARY_POSTFIX;
     case Token.DOT:
     case Token.NEW:
-    case Token.OPENBRACKET: return Precedence.MEMBERACCESS;
+    case Token.OPENBRACKET:
+    case Token.EXCLAMATION: return Precedence.MEMBERACCESS;
   }
   return Precedence.NONE;
 }

src/resolver.ts

@@ -48,7 +48,8 @@ import {
   Expression,
   IntegerLiteralExpression,
   UnaryPrefixExpression,
-  UnaryPostfixExpression
+  UnaryPostfixExpression,
+  AssertionKind
 } from "./ast";
 
 import {
@@ -685,17 +686,29 @@ export class Resolver extends DiagnosticEmitter {
     }
     switch (expression.kind) {
       case NodeKind.ASSERTION: {
+        if ((<AssertionExpression>expression).assertionKind == AssertionKind.NONNULL) {
+          return this.resolveExpression(
+            (<AssertionExpression>expression).expression,
+            contextualFunction,
+            contextualType,
+            reportMode
+          );
+        }
         let type = this.resolveType(
-          (<AssertionExpression>expression).toType,
+          assert((<AssertionExpression>expression).toType),
           contextualFunction.flow.contextualTypeArguments,
           reportMode
         );
         if (!type) return null;
-        let classType = type.classReference;
-        if (!classType) return null;
+        let element: Element | null = type.classReference;
+        if (!element) {
+          let signature = type.signatureReference;
+          if (!signature) return null;
+          element = signature.asFunctionTarget(this.program);
+        }
         this.currentThisExpression = null;
         this.currentElementExpression = null;
-        return classType;
+        return element;
       }
       case NodeKind.UNARYPREFIX: {
         // TODO: overloads
@@ -885,11 +898,7 @@ export class Resolver extends DiagnosticEmitter {
           } else {
             let signature = returnType.signatureReference;
             if (signature) {
-              let functionTarget = signature.cachedFunctionTarget;
-              if (!functionTarget) {
-                functionTarget = new FunctionTarget(this.program, signature);
-                signature.cachedFunctionTarget = functionTarget;
-              }
+              let functionTarget = signature.asFunctionTarget(this.program);
               // reuse resolvedThisExpression (might be property access)
               // reuse resolvedElementExpression (might be element access)
               return functionTarget;

src/types.ts

@@ -526,6 +526,13 @@ export class Signature {
     this.type = Type.u32.asFunction(this);
   }
 
+  asFunctionTarget(program: Program): FunctionTarget {
+    var target = this.cachedFunctionTarget;
+    if (!target) this.cachedFunctionTarget = target = new FunctionTarget(program, this);
+    else assert(target.program == program);
+    return target;
+  }
+
   /** Gets the known or, alternatively, generic parameter name at the specified index. */
   getParameterName(index: i32): string {
     var parameterNames = this.parameterNames;

tests/compiler/nonNullAssertion.optimized.wat

@@ -0,0 +1,88 @@
+(module
+ (type $ii (func (param i32) (result i32)))
+ (type $i (func (result i32)))
+ (type $v (func))
+ (memory $0 0)
+ (table $0 1 anyfunc)
+ (elem (i32.const 0) $null)
+ (global $~lib/allocator/arena/startOffset (mut i32) (i32.const 0))
+ (global $~lib/allocator/arena/offset (mut i32) (i32.const 0))
+ (global $~argc (mut i32) (i32.const 0))
+ (export "memory" (memory $0))
+ (export "table" (table $0))
+ (export "testVar" (func $nonNullAssertion/testVar))
+ (export "testObj" (func $nonNullAssertion/testObj))
+ (export "testProp" (func $nonNullAssertion/testObj))
+ (export "testArr" (func $nonNullAssertion/testArr))
+ (export "testElem" (func $nonNullAssertion/testArr))
+ (export "testAll" (func $nonNullAssertion/testAll))
+ (export "testAll2" (func $nonNullAssertion/testAll))
+ (export "testFn" (func $nonNullAssertion/testFn))
+ (export "testFn2" (func $nonNullAssertion/testFn))
+ (export "testRet" (func $nonNullAssertion/testFn))
+ (export "testObjFn" (func $nonNullAssertion/testObjFn))
+ (export "testObjRet" (func $nonNullAssertion/testObjFn))
+ (start $start)
+ (func $nonNullAssertion/testVar (; 0 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+ )
+ (func $nonNullAssertion/testObj (; 1 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+  i32.load
+ )
+ (func $nonNullAssertion/testArr (; 2 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  get_local $0
+  i32.load
+  tee_local $0
+  i32.load
+  i32.const 2
+  i32.shr_u
+  i32.lt_u
+  if (result i32)
+   get_local $0
+   i32.load offset=8
+  else   
+   unreachable
+  end
+ )
+ (func $nonNullAssertion/testAll (; 3 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  get_local $0
+  i32.load
+  tee_local $0
+  i32.load
+  i32.const 2
+  i32.shr_u
+  i32.lt_u
+  if (result i32)
+   get_local $0
+   i32.load offset=8
+  else   
+   unreachable
+  end
+  i32.load
+ )
+ (func $nonNullAssertion/testFn (; 4 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  set_global $~argc
+  get_local $0
+  call_indirect (type $i)
+ )
+ (func $nonNullAssertion/testObjFn (; 5 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  set_global $~argc
+  get_local $0
+  i32.load offset=4
+  call_indirect (type $i)
+ )
+ (func $start (; 6 ;) (type $v)
+  i32.const 8
+  set_global $~lib/allocator/arena/startOffset
+  get_global $~lib/allocator/arena/startOffset
+  set_global $~lib/allocator/arena/offset
+ )
+ (func $null (; 7 ;) (type $v)
+  nop
+ )
+)

tests/compiler/nonNullAssertion.ts

@@ -0,0 +1,55 @@
+import "allocator/arena";
+
+export function testVar(n: Error | null): Error {
+  return n!;
+}
+
+class Foo {
+  bar: Foo | null;
+  baz: (() => Foo | null) | null;
+}
+
+export function testObj(foo: Foo | null): Foo | null {
+  return foo!.bar;
+}
+
+export function testProp(foo: Foo): Foo {
+  return foo.bar!;
+}
+
+export function testArr(foo: Array<Foo> | null): Foo {
+  return foo![0];
+}
+
+export function testElem(foo: Array<Foo | null>): Foo {
+  return foo[0]!;
+}
+
+export function testAll(foo: Array<Foo | null> | null): Foo {
+  return foo![0]!.bar!;
+}
+
+export function testAll2(foo: Array<Foo | null> | null): Foo {
+  return foo!![0]!!!.bar!!!!;
+}
+
+export function testFn(fn: (() => Foo | null) | null): Foo | null {
+  return fn!();
+}
+
+export function testFn2(fn: (() => Foo | null) | null): Foo | null {
+  var fn2 = fn!;
+  return fn2();
+}
+
+export function testRet(fn: (() => Foo | null) | null): Foo {
+  return fn!()!;
+}
+
+export function testObjFn(foo: Foo): Foo | null {
+  return foo.baz!();
+}
+
+export function testObjRet(foo: Foo): Foo {
+  return foo.baz!()!;
+}

tests/compiler/nonNullAssertion.untouched.wat

@@ -0,0 +1,169 @@
+(module
+ (type $ii (func (param i32) (result i32)))
+ (type $iii (func (param i32 i32) (result i32)))
+ (type $i (func (result i32)))
+ (type $v (func))
+ (memory $0 0)
+ (table $0 1 anyfunc)
+ (elem (i32.const 0) $null)
+ (global $~lib/internal/allocator/AL_BITS i32 (i32.const 3))
+ (global $~lib/internal/allocator/AL_SIZE i32 (i32.const 8))
+ (global $~lib/internal/allocator/AL_MASK i32 (i32.const 7))
+ (global $~lib/internal/allocator/MAX_SIZE_32 i32 (i32.const 1073741824))
+ (global $~lib/allocator/arena/startOffset (mut i32) (i32.const 0))
+ (global $~lib/allocator/arena/offset (mut i32) (i32.const 0))
+ (global $~lib/internal/arraybuffer/HEADER_SIZE i32 (i32.const 8))
+ (global $~argc (mut i32) (i32.const 0))
+ (global $HEAP_BASE i32 (i32.const 8))
+ (export "memory" (memory $0))
+ (export "table" (table $0))
+ (export "testVar" (func $nonNullAssertion/testVar))
+ (export "testObj" (func $nonNullAssertion/testObj))
+ (export "testProp" (func $nonNullAssertion/testProp))
+ (export "testArr" (func $nonNullAssertion/testArr))
+ (export "testElem" (func $nonNullAssertion/testElem))
+ (export "testAll" (func $nonNullAssertion/testAll))
+ (export "testAll2" (func $nonNullAssertion/testAll2))
+ (export "testFn" (func $nonNullAssertion/testFn))
+ (export "testFn2" (func $nonNullAssertion/testFn2))
+ (export "testRet" (func $nonNullAssertion/testRet))
+ (export "testObjFn" (func $nonNullAssertion/testObjFn))
+ (export "testObjRet" (func $nonNullAssertion/testObjRet))
+ (start $start)
+ (func $nonNullAssertion/testVar (; 0 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+ )
+ (func $nonNullAssertion/testObj (; 1 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+  i32.load
+ )
+ (func $nonNullAssertion/testProp (; 2 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+  i32.load
+ )
+ (func $~lib/array/Array<Foo>#__get (; 3 ;) (type $iii) (param $0 i32) (param $1 i32) (result i32)
+  (local $2 i32)
+  (local $3 i32)
+  get_local $0
+  i32.load
+  set_local $2
+  get_local $1
+  get_local $2
+  i32.load
+  i32.const 2
+  i32.shr_u
+  i32.lt_u
+  if (result i32)
+   i32.const 0
+   set_local $3
+   get_local $2
+   get_local $1
+   i32.const 2
+   i32.shl
+   i32.add
+   get_local $3
+   i32.add
+   i32.load offset=8
+  else   
+   unreachable
+  end
+ )
+ (func $nonNullAssertion/testArr (; 4 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+  i32.const 0
+  call $~lib/array/Array<Foo>#__get
+ )
+ (func $~lib/array/Array<Foo | null>#__get (; 5 ;) (type $iii) (param $0 i32) (param $1 i32) (result i32)
+  (local $2 i32)
+  (local $3 i32)
+  get_local $0
+  i32.load
+  set_local $2
+  get_local $1
+  get_local $2
+  i32.load
+  i32.const 2
+  i32.shr_u
+  i32.lt_u
+  if (result i32)
+   i32.const 0
+   set_local $3
+   get_local $2
+   get_local $1
+   i32.const 2
+   i32.shl
+   i32.add
+   get_local $3
+   i32.add
+   i32.load offset=8
+  else   
+   unreachable
+  end
+ )
+ (func $nonNullAssertion/testElem (; 6 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+  i32.const 0
+  call $~lib/array/Array<Foo | null>#__get
+ )
+ (func $nonNullAssertion/testAll (; 7 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+  i32.const 0
+  call $~lib/array/Array<Foo | null>#__get
+  i32.load
+ )
+ (func $nonNullAssertion/testAll2 (; 8 ;) (type $ii) (param $0 i32) (result i32)
+  get_local $0
+  i32.const 0
+  call $~lib/array/Array<Foo | null>#__get
+  i32.load
+ )
+ (func $nonNullAssertion/testFn (; 9 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  set_global $~argc
+  get_local $0
+  call_indirect (type $i)
+ )
+ (func $nonNullAssertion/testFn2 (; 10 ;) (type $ii) (param $0 i32) (result i32)
+  (local $1 i32)
+  get_local $0
+  set_local $1
+  i32.const 0
+  set_global $~argc
+  get_local $1
+  call_indirect (type $i)
+ )
+ (func $nonNullAssertion/testRet (; 11 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  set_global $~argc
+  get_local $0
+  call_indirect (type $i)
+ )
+ (func $nonNullAssertion/testObjFn (; 12 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  set_global $~argc
+  get_local $0
+  i32.load offset=4
+  call_indirect (type $i)
+ )
+ (func $nonNullAssertion/testObjRet (; 13 ;) (type $ii) (param $0 i32) (result i32)
+  i32.const 0
+  set_global $~argc
+  get_local $0
+  i32.load offset=4
+  call_indirect (type $i)
+ )
+ (func $start (; 14 ;) (type $v)
+  get_global $HEAP_BASE
+  get_global $~lib/internal/allocator/AL_MASK
+  i32.add
+  get_global $~lib/internal/allocator/AL_MASK
+  i32.const -1
+  i32.xor
+  i32.and
+  set_global $~lib/allocator/arena/startOffset
+  get_global $~lib/allocator/arena/startOffset
+  set_global $~lib/allocator/arena/offset
+ )
+ (func $null (; 15 ;) (type $v)
+ )
+)

tests/parser/nonNullAssertion.ts

@@ -0,0 +1,8 @@
+foo!;
+foo!!;
+foo!.bar;
+foo.bar!;
+foo![0];
+foo[0]!;
+foo![0]!.bar!;
+foo!![0]!!.bar!!;

tests/parser/nonNullAssertion.ts.fixture.ts

@@ -0,0 +1,8 @@
+foo!;
+foo!!;
+foo!.bar;
+foo.bar!;
+foo![0];
+foo[0]!;
+foo![0]!.bar!;
+foo!![0]!!.bar!!;