Fix managed function expression return

src/compiler.ts

@@ -1141,12 +1141,13 @@ export class Compiler extends DiagnosticEmitter {
       // none of the following can be an arrow function
       assert(!instance.isAny(CommonFlags.CONSTRUCTOR | CommonFlags.GET | CommonFlags.SET));
 
-      let expr = this.compileExpression((<ExpressionStatement>bodyNode).expression, returnType,
-        Constraints.CONV_IMPLICIT
-      );
+      // take special care of properly retaining the returned value
+      let expr = this.compileReturnedExpression((<ExpressionStatement>bodyNode).expression, returnType, Constraints.CONV_IMPLICIT);
+
       if (!stmts) stmts = [ expr ];
       else stmts.push(expr);
-      if (!flow.is(FlowFlags.TERMINATES)) { // TODO: detect if returning an autorelease local?
+
+      if (!flow.is(FlowFlags.TERMINATES)) {
         let indexBefore = stmts.length;
         this.performAutoreleases(flow, stmts);
         this.finishAutoreleases(flow, stmts);
@@ -2215,6 +2216,32 @@ export class Compiler extends DiagnosticEmitter {
     // foo // is possibly null
   }
 
+  /** Compiles an expression that is about to be returned, taking special care of retaining and setting flow states. */
+  compileReturnedExpression(
+    /** Expression to compile. */
+    expression: Expression,
+    /** Return type of the function. */
+    returnType: Type,
+    /** Constraints indicating contextual conditions. */
+    constraints: Constraints = Constraints.NONE
+  ): ExpressionRef {
+    // pretend to retain the expression immediately so the autorelease, if any, is skipped
+    var expr = this.compileExpression(expression, returnType, constraints | Constraints.WILL_RETAIN);
+    var flow = this.currentFlow;
+    if (returnType.isManaged) {
+      // check if that worked, and if it didn't, keep the reference alive
+      if (!this.skippedAutoreleases.has(expr)) {
+        let index = this.tryUndoAutorelease(expr, flow);
+        if (index == -1) expr = this.makeRetain(expr);
+        this.skippedAutoreleases.add(expr);
+      }
+    }
+    // remember return states
+    if (!flow.canOverflow(expr, returnType)) flow.set(FlowFlags.RETURNS_WRAPPED);
+    if (flow.isNonnull(expr, returnType)) flow.set(FlowFlags.RETURNS_NONNULL);
+    return expr;
+  }
+
   compileReturnStatement(
     statement: ReturnStatement,
     isLastInBody: bool
@@ -2239,27 +2266,9 @@ export class Compiler extends DiagnosticEmitter {
       }
       let constraints = Constraints.CONV_IMPLICIT;
       if (flow.actualFunction.is(CommonFlags.MODULE_EXPORT)) constraints |= Constraints.MUST_WRAP;
-      expr = this.compileExpression(valueExpression, returnType, constraints | Constraints.WILL_RETAIN);
-
-      // when returning a local, and it is already retained, skip the final set
-      // of retaining it as the return value and releasing it as a variable
-      if (!this.skippedAutoreleases.has(expr)) {
-        if (returnType.isManaged) {
-          if (getExpressionId(expr) == ExpressionId.LocalGet) {
-            let index = getLocalGetIndex(expr);
-            if (flow.isAnyLocalFlag(index, LocalFlags.ANY_RETAINED)) {
-              flow.unsetLocalFlag(index, LocalFlags.ANY_RETAINED);
-              flow.setLocalFlag(index, LocalFlags.RETURNED);
-              this.skippedAutoreleases.add(expr);
-            }
-          }
-        }
-      }
-
-      // remember return states
-      if (!flow.canOverflow(expr, returnType)) flow.set(FlowFlags.RETURNS_WRAPPED);
-      if (flow.isNonnull(expr, returnType)) flow.set(FlowFlags.RETURNS_NONNULL);
 
+      // take special care of properly retaining the returned value
+      expr = this.compileReturnedExpression(valueExpression, returnType, constraints);
     } else if (returnType != Type.void) {
       this.error(
         DiagnosticCode.Type_0_is_not_assignable_to_type_1,
@@ -2272,9 +2281,6 @@ export class Compiler extends DiagnosticEmitter {
     this.performAutoreleases(flow, stmts);
     this.finishAutoreleases(flow, stmts);
 
-    // Make sure that the return value is retained for the caller
-    if (returnType.isManaged && !this.skippedAutoreleases.has(expr)) expr = this.makeRetain(expr);
-
     if (returnType != Type.void && stmts.length) {
       let temp = flow.getTempLocal(returnType);
       if (flow.isNonnull(expr, returnType)) flow.setLocalFlag(temp.index, LocalFlags.NONNULL);
@@ -6506,24 +6512,32 @@ export class Compiler extends DiagnosticEmitter {
     /** Flow that would autorelease. */
     flow: Flow
   ): i32 {
-    // NOTE: Can't remove the local.tee completely because it's already compiled
-    // and a child of something else. Preventing the final release however makes
-    // it optimize away.
+    // The following assumes that the expression actually belongs to the flow and that
+    // top-level autoreleases are never undone. While that's true, it's not necessary
+    // to check presence in scopedLocals.
     switch (getExpressionId(expr)) {
-      case ExpressionId.LocalSet: { // local.tee(__retain(expr))
+      case ExpressionId.LocalGet: { // local.get(idx)
+        let index = getLocalGetIndex(expr);
+        if (flow.isAnyLocalFlag(index, LocalFlags.ANY_RETAINED)) {
+          flow.unsetLocalFlag(index, LocalFlags.ANY_RETAINED);
+          return index;
+        }
+        break;
+      }
+      case ExpressionId.LocalSet: { // local.tee(idx, expr)
         if (isLocalTee(expr)) {
+          // NOTE: Can't remove the local.tee completely because it's already compiled
+          // and a child of something else. Preventing the final release however makes
+          // it optimize away.
           let index = getLocalSetIndex(expr);
           if (flow.isAnyLocalFlag(index, LocalFlags.ANY_RETAINED)) {
-            // Assumes that the expression actually belongs to the flow and that
-            // top-level autoreleases are never undone. While that's true, it's
-            // not necessary to check presence in scopedLocals.
             flow.unsetLocalFlag(index, LocalFlags.ANY_RETAINED);
             return index;
           }
         }
         break;
       }
-      case ExpressionId.Block: { // { ..., local.tee(__retain(expr)) }
+      case ExpressionId.Block: { // { ..., local.get|tee(...) }
         if (getBlockName(expr) === null) { // must not be a break target
           let count = getBlockChildCount(expr);
           if (count) {

src/flow.ts

@@ -156,17 +156,15 @@ export enum LocalFlags {
   WRITTENTO = 1 << 5,
   /** Local is retained. */
   RETAINED = 1 << 6,
-  /** Local is returned. */
-  RETURNED = 1 << 7,
 
   /** Local is conditionally read from. */
-  CONDITIONALLY_READFROM = 1 << 8,
+  CONDITIONALLY_READFROM = 1 << 7,
   /** Local is conditionally written to. */
-  CONDITIONALLY_WRITTENTO = 1 << 9,
+  CONDITIONALLY_WRITTENTO = 1 << 8,
   /** Local must be conditionally retained. */
-  CONDITIONALLY_RETAINED = 1 << 10,
+  CONDITIONALLY_RETAINED = 1 << 9,
   /** Local is conditionally returned. */
-  CONDITIONALLY_RETURNED = 1 << 11,
+  CONDITIONALLY_RETURNED = 1 << 10,
 
   /** Any categorical flag. */
   ANY_CATEGORICAL = CONSTANT
@@ -175,8 +173,7 @@ export enum LocalFlags {
                   | NONNULL
                   | READFROM
                   | WRITTENTO
-                  | RETAINED
-                  | RETURNED,
+                  | RETAINED,
 
   /** Any conditional flag. */
   ANY_CONDITIONAL = RETAINED
@@ -191,11 +188,7 @@ export enum LocalFlags {
 
   /** Any retained flag. */
   ANY_RETAINED = RETAINED
-               | CONDITIONALLY_RETAINED,
-
-  /** Any returned flag. */
-  ANY_RETURNED = RETURNED
-               | CONDITIONALLY_RETURNED
+               | CONDITIONALLY_RETAINED
 }
 export namespace LocalFlags {
   export function join(left: LocalFlags, right: LocalFlags): LocalFlags {
@@ -585,7 +578,6 @@ export class Flow {
       if (flags & LocalFlags.RETAINED)  this.setLocalFlag(i, LocalFlags.CONDITIONALLY_RETAINED);
       if (flags & LocalFlags.READFROM)  this.setLocalFlag(i, LocalFlags.CONDITIONALLY_READFROM);
       if (flags & LocalFlags.WRITTENTO) this.setLocalFlag(i, LocalFlags.CONDITIONALLY_WRITTENTO);
-      if (flags & LocalFlags.RETURNED)  this.setLocalFlag(i, LocalFlags.CONDITIONALLY_RETURNED);
     }
   }
 

tests/compiler/assert-nonnull.untouched.wat

@@ -273,22 +273,15 @@
  )
  (func $assert-nonnull/testFn (; 14 ;) (param $0 i32) (result i32)
   (local $1 i32)
-  (local $2 i32)
   i32.const 0
   global.set $~lib/argc
   local.get $0
   call_indirect (type $none_=>_i32)
   local.tee $1
-  call $~lib/rt/stub/__retain
-  local.set $2
-  local.get $1
-  call $~lib/rt/stub/__release
-  local.get $2
  )
  (func $assert-nonnull/testFn2 (; 15 ;) (param $0 i32) (result i32)
   (local $1 i32)
   (local $2 i32)
-  (local $3 i32)
   local.get $0
   local.tee $1
   if (result i32)
@@ -302,11 +295,6 @@
   local.get $2
   call_indirect (type $none_=>_i32)
   local.tee $1
-  call $~lib/rt/stub/__retain
-  local.set $3
-  local.get $1
-  call $~lib/rt/stub/__release
-  local.get $3
  )
  (func $assert-nonnull/testRet (; 16 ;) (param $0 i32) (result i32)
   (local $1 i32)
@@ -340,10 +328,7 @@
   i32.load offset=4
   call_indirect (type $none_=>_i32)
   local.tee $1
-  call $~lib/rt/stub/__retain
   local.set $2
-  local.get $1
-  call $~lib/rt/stub/__release
   local.get $0
   call $~lib/rt/stub/__release
   local.get $2

tests/compiler/retain-return.json

@@ -0,0 +1,6 @@
+{
+  "asc_flags": [
+    "--runtime half",
+    "--explicitStart"
+  ]
+}