feat: optimize string literal tostack behavior

String literals are widely used static data, and our GC algorithm does not recycle them (see https://github.com/AssemblyScript/assemblyscript/blob/7e2a62d9615d2ae02b593af87ee4920a3c57b0bd/std/assembly/rt/itcms.ts#L244). Recording the addresses of all string literals and preventing the emission of tostack code can effectively improve performance

Author: HerrCai0907

src/compiler.ts

@@ -59,7 +59,8 @@ import {
   isConstNegZero,
   isConstExpressionNaN,
   ensureType,
-  createType
+  createType,
+  getConstValueInteger
 } from "./module";
 
 import {
@@ -451,6 +452,8 @@ export class Compiler extends DiagnosticEmitter {
   memorySegments: MemorySegment[] = [];
   /** Map of already compiled static string segments. */
   stringSegments: Map<string,MemorySegment> = new Map();
+  /** Set of static GC object offsets. tostack is unnecessary for them. */
+  staticGcObjectOffsets: Map<i32, Set<i32>> = new Map();
   /** Function table being compiled. First elem is blank. */
   functionTable: Function[] = [];
   /** Arguments length helper global. */
@@ -1939,7 +1942,16 @@ export class Compiler extends DiagnosticEmitter {
       stringSegment = this.addRuntimeMemorySegment(buf);
       segments.set(stringValue, stringSegment);
     }
-    return i64_add(stringSegment.offset, i64_new(totalOverhead));
+    let stringOffset = i64_add(stringSegment.offset, i64_new(totalOverhead));
+    let staticGcObjectOffsets = this.staticGcObjectOffsets;
+    if (staticGcObjectOffsets.has(i64_high(stringOffset))) {
+      staticGcObjectOffsets.get(i64_high(stringOffset))!.add(i64_low(stringOffset));
+    } else {
+      let s = new Set<i32>();
+      s.add(i64_low(stringOffset));
+      staticGcObjectOffsets.set(i64_high(stringOffset), s);
+    }
+    return stringOffset;
   }
 
   /** Writes a series of static values of the specified type to a buffer. */
@@ -6754,6 +6766,21 @@ export class Compiler extends DiagnosticEmitter {
     stub.set(CommonFlags.Compiled);
   }
 
+  private needToStack(expr: ExpressionRef): bool {
+    const precomp = this.module.runExpression(expr, ExpressionRunnerFlags.Default);
+    // cannot precompute, so must go to stack
+    if (precomp == 0) return true;
+    const value = getConstValueInteger(precomp, this.options.isWasm64);
+    // zero constant doesn't need to go to stack
+    if (i64_eq(value, i64_zero)) return false;
+    // static GC objects doesn't need to go to stack
+    let staticGcObjectOffsets = this.staticGcObjectOffsets;
+    if (staticGcObjectOffsets.has(i64_high(value))) {
+      if (staticGcObjectOffsets.get(i64_high(value))!.has(i64_low(value))) return false;
+    }
+    return true;
+  }
+
   /** Marks managed call operands for the shadow stack. */
   private operandsTostack(signature: Signature, operands: ExpressionRef[]): void {
     if (!this.options.stackSize) return;
@@ -6763,8 +6790,7 @@ export class Compiler extends DiagnosticEmitter {
     if (thisType) {
       if (thisType.isManaged) {
         let operand = operands[0];
-        let precomp = module.runExpression(operand, ExpressionRunnerFlags.Default);
-        if (!precomp || !isConstZero(precomp)) { // otherwise unnecessary
+        if (this.needToStack(operand)) {
           operands[operandIndex] = module.tostack(operand);
         }
       }
@@ -6777,8 +6803,7 @@ export class Compiler extends DiagnosticEmitter {
       let paramType = parameterTypes[parameterIndex];
       if (paramType.isManaged) {
         let operand = operands[operandIndex];
-        let precomp = module.runExpression(operand, ExpressionRunnerFlags.Default);
-        if (!precomp || !isConstZero(precomp)) { // otherwise unnecessary
+        if (this.needToStack(operand)) {
           operands[operandIndex] = module.tostack(operand);
         }
       }

src/module.ts

@@ -3038,6 +3038,17 @@ export function getConstValueI64High(expr: ExpressionRef): i32 {
   return binaryen._BinaryenConstGetValueI64High(expr);
 }
 
+export function getConstValueInteger(expr: ExpressionRef, isWasm64: bool): i64 {
+    let lo, hi = 0;
+    if (isWasm64) {
+      lo = getConstValueI64Low(expr);
+      hi = getConstValueI64High(expr);
+    } else {
+      lo = getConstValueI32(expr);
+    }
+    return i64_new(lo, hi);
+}
+
 export function getConstValueF32(expr: ExpressionRef): f32 {
   return binaryen._BinaryenConstGetValueF32(expr);
 }

tests/compiler/NonNullable.debug.wat

@@ -357,28 +357,15 @@
   (local $0 i32)
   (local $1 i32)
   global.get $~lib/memory/__stack_pointer
-  i32.const 12
+  i32.const 8
   i32.sub
   global.set $~lib/memory/__stack_pointer
   call $~stack_check
   global.get $~lib/memory/__stack_pointer
   i64.const 0
   i64.store
-  global.get $~lib/memory/__stack_pointer
-  i32.const 0
-  i32.store offset=8
   i32.const 32
-  local.set $1
-  global.get $~lib/memory/__stack_pointer
-  local.get $1
-  i32.store
-  local.get $1
   i32.const 32
-  local.set $1
-  global.get $~lib/memory/__stack_pointer
-  local.get $1
-  i32.store offset=4
-  local.get $1
   call $~lib/string/String.__eq
   i32.eqz
   if
@@ -390,17 +377,7 @@
    unreachable
   end
   i32.const 112
-  local.set $1
-  global.get $~lib/memory/__stack_pointer
-  local.get $1
-  i32.store
-  local.get $1
   i32.const 112
-  local.set $1
-  global.get $~lib/memory/__stack_pointer
-  local.get $1
-  i32.store offset=4
-  local.get $1
   call $~lib/string/String.__eq
   i32.eqz
   if
@@ -412,17 +389,7 @@
    unreachable
   end
   i32.const 144
-  local.set $1
-  global.get $~lib/memory/__stack_pointer
-  local.get $1
-  i32.store
-  local.get $1
   i32.const 144
-  local.set $1
-  global.get $~lib/memory/__stack_pointer
-  local.get $1
-  i32.store offset=4
-  local.get $1
   call $~lib/string/String.__eq
   i32.eqz
   if
@@ -436,7 +403,7 @@
   global.get $~lib/memory/__stack_pointer
   global.get $NonNullable/z
   local.tee $0
-  i32.store offset=8
+  i32.store offset=4
   local.get $0
   if (result i32)
    local.get $0
@@ -462,7 +429,7 @@
   local.get $1
   call $NonNullable/safetyCheck<~lib/string/String|null>
   global.get $~lib/memory/__stack_pointer
-  i32.const 12
+  i32.const 8
   i32.add
   global.set $~lib/memory/__stack_pointer
  )

tests/compiler/NonNullable.release.wat

@@ -22,7 +22,7 @@
  (start $~start)
  (func $~start
   global.get $~lib/memory/__stack_pointer
-  i32.const 12
+  i32.const 8
   i32.sub
   global.set $~lib/memory/__stack_pointer
   block $folding-inner0
@@ -33,15 +33,6 @@
    global.get $~lib/memory/__stack_pointer
    i64.const 0
    i64.store
-   global.get $~lib/memory/__stack_pointer
-   i32.const 0
-   i32.store offset=8
-   global.get $~lib/memory/__stack_pointer
-   i32.const 1056
-   i32.store
-   global.get $~lib/memory/__stack_pointer
-   i32.const 1056
-   i32.store offset=4
    i32.const 1056
    i32.const 1056
    call $~lib/string/String.__eq
@@ -54,12 +45,6 @@
     call $~lib/builtins/abort
     unreachable
    end
-   global.get $~lib/memory/__stack_pointer
-   i32.const 1136
-   i32.store
-   global.get $~lib/memory/__stack_pointer
-   i32.const 1136
-   i32.store offset=4
    i32.const 1136
    i32.const 1136
    call $~lib/string/String.__eq
@@ -72,12 +57,6 @@
     call $~lib/builtins/abort
     unreachable
    end
-   global.get $~lib/memory/__stack_pointer
-   i32.const 1168
-   i32.store
-   global.get $~lib/memory/__stack_pointer
-   i32.const 1168
-   i32.store offset=4
    i32.const 1168
    i32.const 1168
    call $~lib/string/String.__eq
@@ -92,7 +71,7 @@
    end
    global.get $~lib/memory/__stack_pointer
    i32.const 1248
-   i32.store offset=8
+   i32.store offset=4
    global.get $~lib/memory/__stack_pointer
    i32.const 1248
    i32.store
@@ -126,7 +105,7 @@
    i32.add
    global.set $~lib/memory/__stack_pointer
    global.get $~lib/memory/__stack_pointer
-   i32.const 12
+   i32.const 8
    i32.add
    global.set $~lib/memory/__stack_pointer
    return

tests/compiler/bindings/esm.debug.wat

@@ -124,6 +124,27 @@
  (export "functionFunction" (func $export:bindings/esm/functionFunction))
  (func $start:bindings/esm~anonymous|0
  )
+ (func $start:bindings/esm
+  i32.const 128
+  i32.const 1
+  f64.const 42
+  f64.const 0
+  f64.const 0
+  f64.const 0
+  f64.const 0
+  call $~lib/builtins/trace
+  i32.const 160
+  call $~lib/bindings/dom/console.log
+  global.get $~lib/bindings/dom/Math.E
+  call $~lib/bindings/dom/Math.log
+  drop
+  global.get $bindings/esm/immutableGlobal
+  drop
+  global.get $bindings/esm/immutableGlobalNested
+  drop
+  call $bindings/esm/Date_getTimezoneOffset
+  drop
+ )
  (func $bindings/esm/plainFunction (param $a i32) (param $b i32) (result i32)
   local.get $a
   local.get $b
@@ -3081,50 +3102,6 @@
    unreachable
   end
  )
- (func $start:bindings/esm
-  (local $0 i32)
-  global.get $~lib/memory/__stack_pointer
-  i32.const 4
-  i32.sub
-  global.set $~lib/memory/__stack_pointer
-  call $~stack_check
-  global.get $~lib/memory/__stack_pointer
-  i32.const 0
-  i32.store
-  i32.const 128
-  local.set $0
-  global.get $~lib/memory/__stack_pointer
-  local.get $0
-  i32.store
-  local.get $0
-  i32.const 1
-  f64.const 42
-  f64.const 0
-  f64.const 0
-  f64.const 0
-  f64.const 0
-  call $~lib/builtins/trace
-  i32.const 160
-  local.set $0
-  global.get $~lib/memory/__stack_pointer
-  local.get $0
-  i32.store
-  local.get $0
-  call $~lib/bindings/dom/console.log
-  global.get $~lib/bindings/dom/Math.E
-  call $~lib/bindings/dom/Math.log
-  drop
-  global.get $bindings/esm/immutableGlobal
-  drop
-  global.get $bindings/esm/immutableGlobalNested
-  drop
-  call $bindings/esm/Date_getTimezoneOffset
-  drop
-  global.get $~lib/memory/__stack_pointer
-  i32.const 4
-  i32.add
-  global.set $~lib/memory/__stack_pointer
- )
  (func $bindings/esm/bufferFunction (param $a i32) (param $b i32) (result i32)
   (local $aByteLength i32)
   (local $bByteLength i32)