now each handler has an inline and prelude section so exception handlers run properly

src/Examples.java

@@ -18,7 +18,7 @@ public class Examples {
         Jokio.runtime().await().spawn(server());
 
 
-        for(int i = 0; i < 100; i ++){
+        for(int i = 0; i < 200; i ++){
             var builder = new StringBuilder();
             for(int c = 0; c < 4096*16*3; c ++)
                 builder.append((char)((Math.random()*('z'-'a')+'a')));

src/generator/runtime/Frame.java

@@ -9,7 +9,7 @@ public record Frame(FrameTracker.Type[] locals, FrameTracker.Type[] stack) {
 
     @Override
     public String toString() {
-        return "Frame[l =" + Arrays.toString(locals) + ", s = " + Arrays.toString(stack) + "]";
+        return "Frame[label =" + Arrays.toString(locals) + ", s = " + Arrays.toString(stack) + "]";
     }
 
     public void save_locals(StateMachineBuilder smb, CodeBuilder cob, SavedStateTracker sst, int loc_off){

src/generator/runtime/SpecialMethodHandler.java

@@ -1,26 +1,13 @@
 package generator.runtime;
 
 import java.lang.classfile.CodeBuilder;
-import java.lang.classfile.Label;
 
-public abstract class SpecialMethodHandler {
-    public final Label handler_start;
-    public final Label handler_resume;
+public interface SpecialMethodHandler {
 
-    protected SpecialMethodHandler(StateMachineBuilder smb, CodeBuilder cob) {
-        handler_start = cob.newLabel();
-        handler_resume = cob.newLabel();
-    }
-
-    public abstract void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame);
-
-    public boolean removeCall() {
+    void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame);
+    default boolean removeCall() {
         return true;
     }
-
-    public void insertShim(StateMachineBuilder smb, CodeBuilder cob){
-        cob.goto_(handler_start).labelBinding(handler_resume);
-    }
-
-    public abstract ReplacementKind replacementKind();
+    void build_inline(StateMachineBuilder smb, CodeBuilder cob, Frame frame);
+    ReplacementKind replacementKind();
 }

src/generator/runtime/StateBuilder.java

@@ -0,0 +1,31 @@
+package generator.runtime;
+
+import java.lang.classfile.CodeBuilder;
+import java.lang.classfile.Label;
+import java.lang.classfile.instruction.SwitchCase;
+import java.lang.constant.ConstantDescs;
+import java.util.ArrayList;
+
+public class StateBuilder {
+    public record State(Label label, int id){
+        public void bind(CodeBuilder cob){
+            cob.labelBinding(label);
+        }
+
+        public void setState(StateMachineBuilder smb, CodeBuilder cob){
+            cob.aload(0).loadConstant(id).putfield(smb.CD_this, StateMachineBuilder.STATE_NAME, ConstantDescs.CD_int);
+        }
+    }
+
+    private final ArrayList<State> states = new ArrayList<>();
+
+    public State create(CodeBuilder cob){
+        var state = new State(cob.newLabel(), states.size());
+        states.add(state);
+        return state;
+    }
+
+    public void buildSwitch(CodeBuilder cob, Label default_label){
+        cob.lookupswitch(default_label, states.stream().map(l -> SwitchCase.of(l.id, l.label)).toList());
+    }
+}

src/generator/runtime/StateMachineBuilder.java

@@ -10,7 +10,6 @@ import java.lang.constant.ConstantDescs;
 import java.lang.constant.MethodTypeDesc;
 import java.lang.reflect.AccessFlag;
 import java.util.*;
-import java.util.function.BiFunction;
 
 public abstract class StateMachineBuilder {
     public final static String PARAM_PREFIX = "param_";
@@ -36,9 +35,13 @@ public abstract class StateMachineBuilder {
     record LState(String name, ClassDesc cd) {
     }
 
-    protected HashMap<SpecialMethod, BiFunction<StateMachineBuilder, CodeBuilder, SpecialMethodHandler>> smmap = new HashMap<>();
+    public interface SpecialMethodBuilder{
+        SpecialMethodHandler build(StateMachineBuilder smb, CodeBuilder cob, StateBuilder sb);
+    }
 
-    private final ArrayList<SwitchCase> stateSwitchCases = new ArrayList<>();
+    protected HashMap<SpecialMethod, SpecialMethodBuilder> smmap = new HashMap<>();
+
+//    private final ArrayList<SwitchCase> stateSwitchCases = new ArrayList<>();
 
     protected String uniqueName(){
         return sequence+++"";
@@ -108,10 +111,10 @@ public abstract class StateMachineBuilder {
         };
     }
 
-    public int add_state(Label label) {
-        stateSwitchCases.add(SwitchCase.of(stateSwitchCases.size(), label));
-        return stateSwitchCases.size() - 1;
-    }
+//    public int add_state(Label label) {
+//        stateSwitchCases.add(SwitchCase.of(stateSwitchCases.size(), label));
+//        return stateSwitchCases.size() - 1;
+//    }
 
     public void buildSourceMethodShim(CodeBuilder cob){
         cob.new_(CD_this).dup();
@@ -179,43 +182,38 @@ public abstract class StateMachineBuilder {
     }
 
     public void buildStateMachineCode(ClassBuilder clb, CodeBuilder cob, int loc_param_off) {
+        var stateBuilder = new StateBuilder();
         boolean ignore_next_pop = false;
 
         var invalid_state = cob.newLabel();
-        var start_label = cob.newLabel();
-        add_state(start_label);
+        var start_state = stateBuilder.create(cob);
 
         for(var wf : with_frames()){
             if (wf.coe() instanceof InvokeInstruction is){
                 var handler = smmap.get(new SpecialMethod(is.owner().asSymbol(), is.name().stringValue(), is.typeSymbol()));
                 if(handler != null)
-                    handlers.add(handler.apply(this, cob));
+                    handlers.add(handler.build(this, cob, stateBuilder));
             }
         }
 
-        cob.aload(0).getfield(CD_this, STATE_NAME, TypeKind.INT.upperBound()).lookupswitch(invalid_state, stateSwitchCases);
-        var start = cob.startLabel();
-        var end = cob.newLabel();
-        cob.localVariable(0, "this", CD_this, start, end);
+        cob.aload(0).getfield(CD_this, STATE_NAME, TypeKind.INT.upperBound());
+        stateBuilder.buildSwitch(cob, invalid_state);
+
+        cob.localVariable(0, "this", CD_this, cob.startLabel(), cob.endLabel());
 
         {
             int i = 0;
             for (var wf : with_frames()) {
                 if (wf.coe() instanceof InvokeInstruction is) {
                     var h = smmap.get(new SpecialMethod(is.owner().asSymbol(), is.name().stringValue(), is.typeSymbol()));
-                    if(h!=null){
-                        var handler = handlers.get(i++);
-                        cob.labelBinding(handler.handler_start);
-                        handler.buildHandler(this, cob, wf.frame());
-                        cob.goto_(handler.handler_resume);
-                    }
+                    if(h!=null) handlers.get(i++).build_prelude(this, cob, wf.frame());
                 }
             }
         }
 
         SpecialMethodHandler currentHandler = null;
 
-        cob.labelBinding(start_label);
+        start_state.bind(cob);
         for (var wf : with_frames()) {
             var coe = wf.coe();
             var frame = wf.frame();
@@ -227,7 +225,7 @@ public abstract class StateMachineBuilder {
                     }else throw new RuntimeException("Expected Pop Instruction");
                 if (i.opcode() == Opcode.ARETURN){
                     if (currentHandler !=null && currentHandler.replacementKind() == ReplacementKind.ReplacingNextReturn){
-                        currentHandler.insertShim(this, cob);
+                        currentHandler.build_inline(this, cob, frame);
                         currentHandler = null;
                         continue;
                     }
@@ -238,9 +236,9 @@ public abstract class StateMachineBuilder {
                     if(currentHandler!=null)throw new RuntimeException("Multiple method handlers at once not supported");
                     var handler = handlers.removeFirst();
                     if(!handler.removeCall()) cob.with(coe);
-                    if(handler.replacementKind() == ReplacementKind.Immediate) handler.insertShim(this, cob);
+                    if(handler.replacementKind() == ReplacementKind.Immediate) handler.build_inline(this, cob, frame);
                     else if(handler.replacementKind() == ReplacementKind.ImmediateReplacingPop) {
-                        handler.insertShim(this, cob);
+                        handler.build_inline(this, cob, frame);
                         ignore_next_pop = true;
                     }else
                         currentHandler = handler;
@@ -285,7 +283,6 @@ public abstract class StateMachineBuilder {
         cob.new_(ClassDesc.ofDescriptor(IllegalStateException.class.descriptorString())).dup()
                 .invokespecial(ClassDesc.ofDescriptor(IllegalStateException.class.descriptorString()), ConstantDescs.INIT_NAME, ConstantDescs.MTD_void)
                 .athrow();
-        cob.labelBinding(end);
 
         for (var lstate : lstate) {
             clb.withField(lstate.name(), lstate.cd(), ClassFile.ACC_PRIVATE);

src/generator/runtime/future/FutureSMBuilder.java

@@ -17,42 +17,51 @@ public class FutureSMBuilder extends StateMachineBuilder {
     public final static ClassDesc CD_Pending = ClassDesc.ofDescriptor(Future.Pending.class.descriptorString());
 
     public final static MethodTypeDesc MTD_Future_Obj = MethodTypeDesc.of(CD_Future, ConstantDescs.CD_Object);
+    public final static MethodTypeDesc MTD_Future = MethodTypeDesc.of(CD_Future);
     public final static MethodTypeDesc MTD_Object_Waker = MethodTypeDesc.of(ConstantDescs.CD_Object, CD_Waker);
     public final static MethodTypeDesc MTD_Obj = MethodTypeDesc.of(ConstantDescs.CD_Object);
 
     public final static String AWAITING_FIELD_NAME = "awaiting";
 
-    static class AwaitHandler extends SpecialMethodHandler{
-        final int awaiting_state;
-        final Label restore_label;
+    static class AwaitHandler implements SpecialMethodHandler{
+        final StateBuilder.State awaiting;
         final Label save_label;
+        final Label resume_inline;
 
-        public AwaitHandler(StateMachineBuilder smb, CodeBuilder cob) {
-            super(smb, cob);
-            awaiting_state = smb.add_state(restore_label = cob.newLabel());
+        public AwaitHandler(StateMachineBuilder smb, CodeBuilder cob, StateBuilder sb) {
+            awaiting = sb.create(cob);
             save_label = cob.newLabel();
+            resume_inline = cob.newLabel();
         }
 
         public ReplacementKind replacementKind(){return ReplacementKind.Immediate;}
 
         @Override
-        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
-            cob.aload(0).loadConstant(awaiting_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound());
+        public void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+            cob.labelBinding(save_label);
+            var sst = new SavedStateTracker();
+            frame.save_locals(smb, cob, sst,2);
+            cob.storeLocal(TypeKind.REFERENCE, frame.locals().length+2);
+            frame.save_stack(smb, cob, sst,1);
+            cob.loadLocal(TypeKind.REFERENCE, frame.locals().length+2);
+            cob.areturn();
+            awaiting.bind(cob);
+            sst.restore_all(smb, cob);
+            cob.goto_(resume_inline);
+        }
+
+        @Override
+        public void build_inline(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+            // [... Future]
             var start = cob.newBoundLabel();
-            cob.dup()
-                    .aload(1).invokeinterface(CD_Future, "poll", MTD_Object_Waker).dup()
-                    .instanceOf(CD_Pending);
+            cob.dup().aload(1).invokeinterface(CD_Future, "poll", MTD_Object_Waker).dup().instanceOf(CD_Pending);
+            // [... Future Polled is_pending]
             cob.ifThenElse(bcb -> {
+                awaiting.setState(smb, cob);
+                // [... Future Polled]
                 bcb.swap().aload(0).swap().putfield(smb.CD_this, AWAITING_FIELD_NAME, CD_Future);
-
-
-                var sst = new SavedStateTracker();
-                frame.save_locals(smb, cob, sst,2);
-                bcb.storeLocal(TypeKind.REFERENCE, frame.locals().length+2);
-                frame.save_stack(smb, cob, sst,1);
-                bcb.loadLocal(TypeKind.REFERENCE, frame.locals().length+2);
-                bcb.areturn().labelBinding(restore_label);
-                sst.restore_all(smb, cob);
+                // [... Polled]
+                cob.goto_(save_label).labelBinding(resume_inline);
 
                 bcb.aload(0).getfield(smb.CD_this, AWAITING_FIELD_NAME, CD_Future);
                 bcb.aload(0).aconst_null().putfield(smb.CD_this, AWAITING_FIELD_NAME, CD_Future);
@@ -60,39 +69,52 @@ public class FutureSMBuilder extends StateMachineBuilder {
             }, bcb -> {
                 bcb.swap().pop();
             });
-
+            // [... Polled]
         }
     }
 
-    static class YieldHandler extends SpecialMethodHandler {
-        final int resume_state;
-        final Label resume_label;
+    static class YieldHandler implements SpecialMethodHandler {
+        final StateBuilder.State resume;
+        final Label save_ret;
+        final Label end;
 
-        public YieldHandler(StateMachineBuilder smb, CodeBuilder cob) {
-            super(smb, cob);
-            resume_state = smb.add_state(resume_label = cob.newLabel());
+        public YieldHandler(StateMachineBuilder smb, CodeBuilder cob, StateBuilder sb) {
+            resume = sb.create(cob);
+            save_ret = cob.newLabel();
+            end = cob.newLabel();
         }
 
         public ReplacementKind replacementKind(){return ReplacementKind.Immediate;}
 
         @Override
-        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+        public void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+            cob.labelBinding(save_ret);
             var saved = frame.save(smb, cob, 2, 0);
-            cob.aload(0).loadConstant(resume_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
-                        .getstatic(CD_Pending, "INSTANCE", CD_Pending)
-                        .areturn().labelBinding(resume_label);
+
+            resume.setState(smb, cob);
+            cob.getstatic(CD_Pending, "INSTANCE", CD_Pending).areturn();
+
+            resume.bind(cob);
             saved.restore_all(smb, cob);
-        }
-    }
-
-    static class WakerHandler extends SpecialMethodHandler{
-
-        protected WakerHandler(StateMachineBuilder smb, CodeBuilder cob) {
-            super(smb, cob);
+            cob.goto_(end);
         }
 
         @Override
-        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+        public void build_inline(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+            cob.goto_(save_ret);
+            cob.labelBinding(end);
+        }
+    }
+
+    static class WakerHandler implements SpecialMethodHandler{
+
+        protected WakerHandler(StateMachineBuilder smb, CodeBuilder cob, StateBuilder sb) {}
+
+        @Override
+        public void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {}
+
+        @Override
+        public void build_inline(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
             cob.aload(1);
         }
 
@@ -102,21 +124,36 @@ public class FutureSMBuilder extends StateMachineBuilder {
         }
     }
 
-    static class RetHandler extends SpecialMethodHandler{
-
-
-        protected RetHandler(StateMachineBuilder smb, CodeBuilder cob) {
-            super(smb, cob);
-        }
+    static class RetHandler implements SpecialMethodHandler{
+        protected RetHandler(StateMachineBuilder smb, CodeBuilder cob, StateBuilder sb) {}
 
         public ReplacementKind replacementKind(){return ReplacementKind.ReplacingNextReturn;}
 
         @Override
-        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+        public void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {}
+
+        @Override
+        public void build_inline(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
             cob.aload(0).loadConstant(-1).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound()).areturn();
         }
     }
 
+    static class RetVoidHandler implements SpecialMethodHandler{
+        protected RetVoidHandler(StateMachineBuilder smb, CodeBuilder cob, StateBuilder sb) {}
+
+        public ReplacementKind replacementKind(){return ReplacementKind.ReplacingNextReturn;}
+
+        @Override
+        public void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+            cob.aload(0).loadConstant(-1).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound()).aconst_null().areturn();
+        }
+
+        @Override
+        public void build_inline(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+
+        }
+    }
+
     @Override
     protected String uniqueName() {
         return "Fut";
@@ -143,6 +180,7 @@ public class FutureSMBuilder extends StateMachineBuilder {
         super(src_clm, src_mem, src_com);
         smmap.put(new SpecialMethod(CD_Future, "await", MTD_Obj), AwaitHandler::new);
         smmap.put(new SpecialMethod(CD_Future, "ret", MTD_Future_Obj), RetHandler::new);
+        smmap.put(new SpecialMethod(CD_Future, "ret", MTD_Future), RetVoidHandler::new);
         smmap.put(new SpecialMethod(CD_Future, "yield", ConstantDescs.MTD_void), YieldHandler::new);
         smmap.put(new SpecialMethod(CD_Waker, "waker", MethodTypeDesc.of(CD_Waker)), WakerHandler::new);
     }

src/generator/runtime/gen/GenSMBuilder.java

@@ -20,64 +20,64 @@ public class GenSMBuilder extends StateMachineBuilder {
     public final static MethodTypeDesc MTD_Gen = MethodTypeDesc.of(CD_Gen);
     public final static MethodTypeDesc MTD_Obj = MethodTypeDesc.of(ConstantDescs.CD_Object);
 
-    static class YieldHandler extends SpecialMethodHandler {
-        final int resume_state;
-        final Label resume_label;
-        final boolean is_void;
-
-        public YieldHandler(StateMachineBuilder smb, CodeBuilder cob, boolean is_void) {
-            super(smb, cob);
-            resume_state = smb.add_state(resume_label = cob.newLabel());
-            this.is_void = is_void;
-        }
-
-        public ReplacementKind replacementKind(){return ReplacementKind.ImmediateReplacingPop;}
-
-        @Override
-        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
-            if(is_void)cob.aconst_null();
-
-//            smb.lt.savingLocals(smb.CD_this, cob, () -> {
-//                cob.aload(0).loadConstant(resume_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
-//                        .new_(CD_Yield)
-//                        .dup_x1()
-//                        .swap()
-//                        .invokespecial(CD_Yield, ConstantDescs.INIT_NAME, MethodTypeDesc.of(ConstantDescs.CD_void, ConstantDescs.CD_Object))
-//                        .areturn();
-//                cob.labelBinding(resume_label);
-//            });
-        }
-    }
-
-    static class RetHandler extends SpecialMethodHandler {
-        final boolean is_void;
-
-        public RetHandler(StateMachineBuilder smb, CodeBuilder cob, boolean is_void) {
-            super(smb, cob);
-            this.is_void = is_void;
-        }
-
-        public ReplacementKind replacementKind(){return ReplacementKind.ReplacingNextReturn;}
-
-        @Override
-        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
-            if(is_void)cob.aconst_null();
-
-            cob.aload(0).loadConstant(-1).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
-                    .new_(CD_Ret)
-                    .dup_x1()
-                    .swap()
-                    .invokespecial(CD_Ret, ConstantDescs.INIT_NAME, MethodTypeDesc.of(ConstantDescs.CD_void, ConstantDescs.CD_Object))
-                    .areturn();
-        }
-    }
+//    static class YieldHandler extends SpecialMethodHandler {
+//        final int resume_state;
+//        final Label resume_label;
+//        final boolean is_void;
+//
+//        public YieldHandler(StateMachineBuilder smb, CodeBuilder cob, boolean is_void) {
+//            super(smb, cob);
+//            resume_state = smb.add_state(resume_label = cob.newLabel());
+//            this.is_void = is_void;
+//        }
+//
+//        public ReplacementKind replacementKind(){return ReplacementKind.ImmediateReplacingPop;}
+//
+//        @Override
+//        public void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+//            if(is_void)cob.aconst_null();
+//
+////            smb.lt.savingLocals(smb.CD_this, cob, () -> {
+////                cob.aload(0).loadConstant(resume_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
+////                        .new_(CD_Yield)
+////                        .dup_x1()
+////                        .swap()
+////                        .invokespecial(CD_Yield, ConstantDescs.INIT_NAME, MethodTypeDesc.of(ConstantDescs.CD_void, ConstantDescs.CD_Object))
+////                        .areturn();
+////                cob.labelBinding(resume_label);
+////            });
+//        }
+//    }
+//
+//    static class RetHandler extends SpecialMethodHandler {
+//        final boolean is_void;
+//
+//        public RetHandler(StateMachineBuilder smb, CodeBuilder cob, boolean is_void) {
+//            super(smb, cob);
+//            this.is_void = is_void;
+//        }
+//
+//        public ReplacementKind replacementKind(){return ReplacementKind.ReplacingNextReturn;}
+//
+//        @Override
+//        public void build_prelude(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
+//            if(is_void)cob.aconst_null();
+//
+//            cob.aload(0).loadConstant(-1).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
+//                    .new_(CD_Ret)
+//                    .dup_x1()
+//                    .swap()
+//                    .invokespecial(CD_Ret, ConstantDescs.INIT_NAME, MethodTypeDesc.of(ConstantDescs.CD_void, ConstantDescs.CD_Object))
+//                    .areturn();
+//        }
+//    }
 
     public GenSMBuilder(ClassModel src_clm, MethodModel src_mem, CodeModel src_com) {
         super(src_clm, src_mem, src_com);
-        smmap.put(new SpecialMethod(CD_Gen, "yield", MTD_Gen_Obj),(smb, cob) -> new YieldHandler(smb, cob, false));
-        smmap.put(new SpecialMethod(CD_Gen, "yield", MTD_Gen),(smb, cob) -> new YieldHandler(smb, cob, true));
-        smmap.put(new SpecialMethod(CD_Gen, "ret", MTD_Gen_Obj),(smb, cob) -> new RetHandler(smb, cob, false));
-        smmap.put(new SpecialMethod(CD_Gen, "ret", MTD_Gen),(smb, cob) -> new RetHandler(smb, cob, true));
+//        smmap.put(new SpecialMethod(CD_Gen, "yield", MTD_Gen_Obj),(smb, cob) -> new YieldHandler(smb, cob, false));
+//        smmap.put(new SpecialMethod(CD_Gen, "yield", MTD_Gen),(smb, cob) -> new YieldHandler(smb, cob, true));
+//        smmap.put(new SpecialMethod(CD_Gen, "ret", MTD_Gen_Obj),(smb, cob) -> new RetHandler(smb, cob, false));
+//        smmap.put(new SpecialMethod(CD_Gen, "ret", MTD_Gen),(smb, cob) -> new RetHandler(smb, cob, true));
     }
 
     @Override