Replace EndOfBlock with check against for stack index

WasmLib/Decompilation/Intermediate/BranchInstruction.cs

@@ -32,7 +32,6 @@ namespace WasmLib.Decompilation.Intermediate
             switch (Kind) {
                 case BranchKind.Normal:
                     context.WriteFull($"BRANCH {Label}");
-                    context.EndOfBlock = true;
                     break;
                 case BranchKind.Conditional:
                     var condition = context.Pop();
@@ -44,7 +43,6 @@ namespace WasmLib.Decompilation.Intermediate
                     var index = context.Pop();
                     Debug.Assert(index.Type == ValueKind.I32);
                     context.WriteFull($"BRANCH_TABLE {{{string.Join(", ", Labels)}}}[{index}] ?? {Label}");
-                    context.EndOfBlock = true;
                     break;
             }
         }

WasmLib/Decompilation/Intermediate/ControlBlockInstruction.cs

@@ -59,32 +59,22 @@ namespace WasmLib.Decompilation.Intermediate
             
             context.WriteFull("}");
 
-            void HandleBlock(ref IntermediateContext contextPassed, IReadOnlyList<IntermediateInstruction> block, bool hasReturnPassed)
+            void HandleBlock(ref IntermediateContext context2, IReadOnlyList<IntermediateInstruction> block, bool hasReturn2)
             {
-                contextPassed.Indent();
-
-                // NOTE: could be optimized if blocks don't pop from stack
-                // TODO: check this by having a Stack<int> of stack sizes in context, and not allowing to go under it
-                var stackBackup = new Stack<Variable>(contextPassed.Stack);
+                context2.EnterBlock();
             
                 foreach (IntermediateInstruction instruction in block) {
-                    instruction.Handle(ref contextPassed);
+                    instruction.Handle(ref context2);
                 }
 
-                if (!contextPassed.EndOfBlock) {
-                    if (hasReturnPassed) {
-                        var popped = contextPassed.Pop();
-                        Debug.Assert(popped.Type == ValueKind);
-                        contextPassed.WriteFull($"block_return {popped}");
-                    }
-                }
-                else {
-                    // block ended on unconditional branch instruction, restore stack to what it was before
-                    contextPassed.RestoreStack(stackBackup);
+                // if stack has values left on it, and we expect a return value
+                if (context2.StackIndices.Peek() != context2.Stack.Count && hasReturn2) {
+                    var popped = context2.Pop();
+                    Debug.Assert(popped.Type == ValueKind);
+                    context2.WriteFull($"block_return {popped}");
                 }
 
-                contextPassed.DeIndent();
-                contextPassed.EndOfBlock = false;
+                context2.ExitBlock();
             }
         }
 

WasmLib/Decompilation/Intermediate/IntermediateContext.cs

@@ -16,9 +16,8 @@ namespace WasmLib.Decompilation.Intermediate
         public WasmFile WasmFile { get; }
         private readonly StreamWriter streamWriter;
         
-        public Stack<Variable> Stack { get; private set; }
-        public bool EndOfBlock { get; set; }
-
+        public Stack<Variable> Stack { get; }
+        public Stack<int> StackIndices { get; }
         private uint varCount;
 
         public IntermediateContext(FunctionBody function, FunctionSignature signature, WasmFile wasmFile, StreamWriter writer)
@@ -32,24 +31,26 @@ namespace WasmLib.Decompilation.Intermediate
                 .Where(x => x.Kind == ImportKind.GlobalType)
                 .Select(x =>
                     x.GlobalType ??
-                    throw new Exception("Import.GlobalType had no value, but Import.Kind was GlobalType"));
+                    throw new Exception($"{nameof(Import.GlobalType)} had no value, but {nameof(Import.Kind)} was {nameof(ImportKind.GlobalType)}"));
             var globals = wasmFile.Globals.Select(x => x.GlobalType); 
             Globals = importGlobals.Concat(globals).Select(x => x.ValueKind).ToList();
             streamWriter = writer;
             Stack = new Stack<Variable>();
-            EndOfBlock = false;
+            StackIndices = new Stack<int>();
+            StackIndices.Push(0);
             varCount = 0;
         }
 
         public Variable Peek()
         {
-            Debug.Assert(Stack.TryPeek(out _), "Tried to peek a value from an empty stack");
+            Debug.Assert(Stack.Any(), "Tried to peek a value from an empty stack");
             return Stack.Peek();
         }
 
         public Variable Pop()
         {
-            Debug.Assert(Stack.TryPeek(out _), "Tried to pop a value from an empty stack");
+            Debug.Assert(Stack.Any(), "Tried to pop a value from an empty stack");
+            Debug.Assert(Stack.Count > StackIndices.Peek(), $"Pop causes stack size to becomes less than {StackIndices.Count}, which is the minimum for this block");
             return Stack.Pop();
         }
 
@@ -60,10 +61,20 @@ namespace WasmLib.Decompilation.Intermediate
             return variable;
         }
 
-        public void RestoreStack(Stack<Variable> newStack)
+        public void EnterBlock()
         {
-            Debug.Assert(EndOfBlock, "Tried to restore a stack but EndOfBlock was not set");
-            Stack = newStack;
+            Indent();
+            StackIndices.Push(Stack.Count);
+        }
+
+        public void ExitBlock()
+        {
+            DeIndent();
+            int previousStackSize = StackIndices.Pop();
+
+            while (Stack.Count > previousStackSize) {
+                Stack.Pop();
+            }
         }
 
         public ValueKind GetLocalType(uint i) => Locals[(int)i];

WasmLib/Decompilation/Intermediate/ReturnInstruction.cs

@@ -17,8 +17,6 @@ namespace WasmLib.Decompilation.Intermediate
             
             // TODO: check if this assert is valid
             // Debug.Assert(context.Stack.Count == 0, "Wrote return instruction while stack was not empty");
-            
-            context.EndOfBlock = true;
         }
     }
 }

WasmLib/Decompilation/Intermediate/VariableInstruction.cs

@@ -49,12 +49,12 @@ namespace WasmLib.Decompilation.Intermediate
                 }
                 else if (Action == ActionKind.Set) {
                     var popped = context.Pop();
-                    Debug.Assert(local == popped.Type);
+                    Debug.Assert(local == popped.Type, $"popped {popped} ({popped.Type}), expected {local}");
                     context.WriteFull($"local[{Index}] = {popped}");
                 }
                 else if (Action == ActionKind.Tee) {
                     var peeked = context.Stack.Peek();
-                    Debug.Assert(local == peeked.Type);
+                    Debug.Assert(local == peeked.Type, $"peeked {peeked} ({peeked.Type}), expected {local}");
                     context.WriteFull($"local[{Index}] = {peeked}");
                 }
             }
@@ -69,7 +69,7 @@ namespace WasmLib.Decompilation.Intermediate
                 else if (Action == ActionKind.Set) {
                     // NOTE: could check for mutability of global
                     var popped = context.Pop();
-                    Debug.Assert(global == popped.Type);
+                    Debug.Assert(global == popped.Type, $"popped {popped} ({popped.Type}), expected {global}");
                     context.WriteFull($"global[{Index}] = {popped}");
                 }
             }

WasmLib/Decompilation/IntermediateRepresentationDecompiler.cs

@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.IO;
 using WasmLib.Decompilation.Intermediate;
 using WasmLib.FileFormat;
@@ -31,7 +30,6 @@ namespace WasmLib.Decompilation
             // write all IR while simulating the stack
             foreach (IntermediateInstruction instruction in instructions) {
                 instruction.Handle(ref context);
-                Debug.Assert(!context.EndOfBlock, "Encountered EndOfBlock instruction when not in block");
             }
             
             // write return value, if needed

WasmTool/Program.cs

@@ -37,7 +37,7 @@ namespace WasmTool
             using var w = new StreamWriter(fs);
             IDecompiler dec = new IntermediateRepresentationDecompiler(wasmFile);
 
-            for (int i = 0; i < Math.Min(wasmFile.FunctionBodies.Length, 520); i++) {
+            for (int i = 0; i < Math.Min(wasmFile.FunctionBodies.Length, 546); i++) {
                 Debug.WriteLine($"Decompiling function {i} (0x{i:X})");
                 dec.DecompileFunction(w, i);
             }