precomputed frames, better local state saving

src/Examples.java

@@ -1,7 +1,10 @@
 import async_example.Delay;
 import async_example.Jokio;
+import async_example.Socket;
 import generator.future.Future;
 
+import java.net.InetSocketAddress;
+
 public class Examples {
     //    public static Gen<String, Void> parse(String str){
 //        for(var gen = parse(str); gen.next() instanceof Gen.Yield(var item);){
@@ -76,8 +79,8 @@ public class Examples {
         var result = awaitTest2(number).await();
         var rt = Jokio.runtime().await();
         rt.spawn(awaitTest2(5000));
-//        closing(100);
+
-//        closing(10).await();
+        closing(5000).await();
         return Future.ret("Result: " + result);
     }
 

src/async_example/Socket.java

@@ -0,0 +1,88 @@
+package async_example;
+
+import generator.future.Future;
+import generator.future.Waker;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.channels.SelectionKey;
+import java.nio.channels.Selector;
+import java.nio.channels.SocketChannel;
+
+public class Socket implements AutoCloseable{
+
+    private final static Selector selector;
+
+    static{
+        try {
+            selector = Selector.open();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+
+        new Thread(() -> {
+            while(!Thread.currentThread().isInterrupted()){
+                try{
+                    selector.select();
+                    var keys = selector.selectedKeys().iterator();
+
+                    while (keys.hasNext()) {
+                        SelectionKey key = keys.next();
+                        keys.remove();
+
+                        if (!key.isValid()) {
+
+                        }else if(key.isAcceptable()){
+
+                        }else if(key.isConnectable()){
+
+                        }else if(key.isReadable()){
+
+                        }else if(key.isWritable()){
+
+                        }
+                    }
+                }catch (Exception e){
+                    e.printStackTrace();
+                }
+            }
+        }).start();
+    }
+
+    private final SocketChannel socket;
+
+    private Socket(SocketChannel sc){
+        this.socket = sc;
+    }
+
+    public static Future<Socket> bind(InetSocketAddress inet) throws IOException{
+        var socket = SocketChannel.open();
+        return new Future<>() {
+            @Override
+            public Socket poll(Waker waker) {
+                try{
+                    socket.configureBlocking(false);
+                    socket.socket().bind(new InetSocketAddress("localhost", 8080));
+                    socket.register(selector, SelectionKey.OP_ACCEPT | SelectionKey.OP_CONNECT | SelectionKey.OP_READ | SelectionKey.OP_WRITE);
+                    return new Socket(socket);
+                }catch (Exception e){
+                    throw new RuntimeException(e);
+                }
+            }
+
+            @Override
+            public void cancel() {
+                try {
+                    socket.close();
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        };
+    }
+
+    @Override
+    public void close() throws Exception {
+        socket.close();
+    }
+}

src/generator/runtime/Frame.java

@@ -0,0 +1,27 @@
+package generator.runtime;
+
+import java.lang.classfile.CodeBuilder;
+import java.lang.classfile.TypeKind;
+import java.util.Arrays;
+
+public record Frame(FrameTracker.Type[] locals, FrameTracker.Type[] stack) {
+
+    @Override
+    public String toString() {
+        return "Frame[l =" + Arrays.toString(locals) + ", s = " + Arrays.toString(stack) + "]";
+    }
+
+    public SavedStateTracker save(StateMachineBuilder smb, CodeBuilder cob, int loc_off, int stack_off) {
+        var sst = new SavedStateTracker();
+        int slot = 0;
+        for (var entry : locals) {
+            slot++;
+            if (slot <= loc_off) continue;
+            if (entry.isCategory2_2nd()) continue;
+
+            sst.save_local(smb, cob, entry.toCD(), slot - smb.paramSlotOff + loc_off - 1);
+        }
+
+        return sst;
+    }
+}

src/generator/runtime/FrameTracker.java

@@ -11,13 +11,9 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Objects;
 
-import static generator.runtime.StateMachineBuilder.LOCAL_PREFIX;
 import static java.lang.constant.ConstantDescs.*;
 
-public class LocalTracker {
+public class FrameTracker {
-
-    record LocalStore(String name, ClassDesc cd) {
-    }
 
     private static final int ITEM_TOP = 0,
             ITEM_INTEGER = 1,
@@ -35,7 +31,14 @@ public class LocalTracker {
             ITEM_LONG_2ND = 13,
             ITEM_DOUBLE_2ND = 14;
 
-    private static record Type(int tag, ClassDesc sym, int bci) {
+    public Type[] locals() {
+        return locals.toArray(Type[]::new);
+    }
+    public Type[] stack() {
+        return stack.toArray(Type[]::new);
+    }
+
+    public record Type(int tag, ClassDesc sym, int bci) {
 
         //singleton types
         static final Type TOP_TYPE = simpleType(ITEM_TOP),
@@ -191,27 +194,19 @@ public class LocalTracker {
 
     final ArrayList<Type> stack = new ArrayList<>();
     final ArrayList<Type> locals = new ArrayList<>();
-//    HashMap<Integer, TypeKind> localVarTypes = new HashMap<>();
-//    StackMapFrameInfo currentFrame;
 
-    HashMap<Integer, ClassDesc> parameter_map = new HashMap<>();
     HashMap<Label, StackMapFrameInfo> stackMapFrames = new HashMap<>();
-    ArrayList<LocalStore> localStore = new ArrayList<>();
 
-    final int local_param_off;
     final StateMachineBuilder smb;
 
 
-    LocalTracker(StateMachineBuilder smb, CodeModel com, int local_param_off) {
+    FrameTracker(StateMachineBuilder smb, CodeModel com) {
-        this.local_param_off = local_param_off;
         this.smb = smb;
         int offset = 0;
 
 
 
         for (var param : smb.params) {
-            parameter_map.put(offset, param);
-
             if(param == CD_long){
                 setLocal2(offset, Type.DOUBLE_TYPE, Type.DOUBLE2_TYPE);
             }else if(param == CD_double){
@@ -246,7 +241,7 @@ public class LocalTracker {
         }
     }
 
-    LocalTracker pushStack(ClassDesc desc) {
+    FrameTracker pushStack(ClassDesc desc) {
         if (desc == CD_long)   return pushStack(Type.LONG_TYPE, Type.LONG2_TYPE);
         if (desc == CD_double) return pushStack(Type.DOUBLE_TYPE, Type.DOUBLE2_TYPE);
         return desc == CD_void ? this
@@ -256,12 +251,12 @@ public class LocalTracker {
                         : Type.referenceType(desc));
     }
 
-    LocalTracker pushStack(Type type) {
+    FrameTracker pushStack(Type type) {
         stack.add(type);
         return this;
     }
 
-    LocalTracker pushStack(Type type1, Type type2) {
+    FrameTracker pushStack(Type type1, Type type2) {
         stack.add(type1);
         stack.add(type2);
         return this;
@@ -271,7 +266,7 @@ public class LocalTracker {
         return stack.removeLast();
     }
 
-    LocalTracker decStack(int size) {
+    FrameTracker decStack(int size) {
         for(int i = 0; i < size; i ++)stack.removeLast();
         return this;
     }
@@ -445,34 +440,6 @@ public class LocalTracker {
             case CustomAttribute<?> _ -> {}
             case StackMapTableAttribute _ -> {}
         }
-        if(ce instanceof Instruction)
-            System.out.println("loc: " + locals + " stack: " + stack);
-    }
-
-
-    //Tries its best to reuse old saved locals field slots, only reuses if types exactly match
-    public void savingLocals(ClassDesc cd, CodeBuilder cob, Runnable run) {
-        record Saved(int slot, String name, ClassDesc cd) {
-        }
-
-        var saved = new ArrayList<Saved>();
-        currentLocals((slot, tk, desc) -> {
-            String name = LOCAL_PREFIX + localStore.size();
-            localStore.add(new LocalStore(name, desc));
-            saved.add(new Saved(slot, name, desc));
-            cob.aload(0).loadLocal(tk, slot).putfield(cd, name, desc);
-        });
-        run.run();
-
-        for (var save : saved) {
-            cob.aload(0).getfield(cd, save.name, save.cd).storeLocal(TypeKind.from(save.cd), save.slot);
-        }
-    }
-
-    public void createLocalStoreFields(ClassBuilder clb) {
-        for (var local : localStore) {
-            clb.withField(local.name, local.cd, ClassFile.ACC_PRIVATE);
-        }
     }
 
     public void encounterLabel(Label l) {
@@ -486,23 +453,4 @@ public class LocalTracker {
                 locals.add(Type.verificationType(local));
         }
     }
-
-    public ClassDesc paramType(int slot) {
-        return parameter_map.get(slot);
-    }
-
-    public interface LocalConsumer {
-        void consume(int slot, TypeKind tk, ClassDesc desc);
-    }
-
-    public void currentLocals(LocalConsumer consumer) {
-        int slot = 0;
-        for (var entry : locals) {
-            slot++;
-            if (slot <= local_param_off) continue;
-            if(entry.isCategory2_2nd())continue;
-            consumer.consume(slot - smb.paramSlotOff + local_param_off - 1, TypeKind.from(entry.toCD()), entry.toCD());
-        }
-
-    }
 }

src/generator/runtime/GeneratorClassLoader.java

@@ -63,11 +63,19 @@ public class GeneratorClassLoader extends ClassLoader {
         return ClassFile.of(ClassFile.AttributesProcessingOption.PASS_ALL_ATTRIBUTES, ClassFile.StackMapsOption.STACK_MAPS_WHEN_REQUIRED).build(clm.thisClass().asSymbol(), cb -> {
             for (var ce : clm) {
                 if (ce instanceof MethodModel mem && !isGen && !isFuture) {
-                    StateMachineBuilder builder = null;
+                    StateMachineBuilder builder;
                     if(mem.methodTypeSymbol().returnType().descriptorString().equals(Gen.class.descriptorString())){
-                        builder = generatorMethod(cb, mem, clm);
+                        builder = new GenSMBuilder(clm, mem, mem.code().get());
                     }else if(mem.methodTypeSymbol().returnType().descriptorString().equals(Future.class.descriptorString())){
-                        builder = futureMethod(cb, mem, clm);
+                        builder = new FutureSMBuilder(clm, mem, mem.code().get());
+                    }else{
+                        builder= null;
+                    }
+                    if(builder!=null&&builder.hasAnyHandlers()){
+                        add(builder.CD_this.displayName(), builder.buildStateMachine());
+                        cb.withMethod(mem.methodName(), mem.methodType(), mem.flags().flagsMask(), mb -> {
+                            mb.withCode(builder::buildSourceMethodShim);
+                        });
                     }else{
                         cb.with(mem);
                     }
@@ -88,33 +96,4 @@ public class GeneratorClassLoader extends ClassLoader {
                 cb.with(NestMembersAttribute.ofSymbols(nestMem));
         });
     }
-
-    private StateMachineBuilder generatorMethod(ClassBuilder cb, MethodModel src_mem, ClassModel src_clm) {
-        var com = src_mem.code().get();
-        var smb = new GenSMBuilder(src_clm, src_mem, com);
-        add(smb.CD_this.displayName(), smb.buildStateMachine());
-        cb.withMethod(src_mem.methodName(), src_mem.methodType(), src_mem.flags().flagsMask(), mb -> {
-            mb.withCode(smb::buildSourceMethodShim);
-        });
-        return smb;
-    }
-
-    private StateMachineBuilder futureMethod(ClassBuilder cb, MethodModel src_mem, ClassModel src_clm) {
-        var com = src_mem.code().get();
-        var smb = new FutureSMBuilder(src_clm, src_mem, com);
-        try{
-            add(smb.CD_this.displayName(), smb.buildStateMachine());
-
-            cb.withMethod(src_mem.methodName(), src_mem.methodType(), src_mem.flags().flagsMask(), mb -> {
-                mb.withCode(smb::buildSourceMethodShim);
-            });
-        }catch (Exception e){
-            e.printStackTrace();
-            cb.withMethod(src_mem.methodName(), src_mem.methodType(), src_mem.flags().flagsMask(), mb -> {
-                mb.with(com);
-            });
-            return smb;
-        }
-        return smb;
-    }
 }

src/generator/runtime/SavedStateTracker.java

@@ -0,0 +1,58 @@
+package generator.runtime;
+
+import java.lang.classfile.CodeBuilder;
+import java.lang.classfile.TypeKind;
+import java.lang.constant.ClassDesc;
+import java.util.ArrayList;
+
+public class SavedStateTracker {
+    private final ArrayList<SavedState> saved = new ArrayList<>();
+    public sealed interface SavedState{}
+    public record StackState(String name, ClassDesc desc) implements SavedState{
+
+    }
+    public record LocalState(String name, ClassDesc desc, int slot) implements SavedState{
+
+    }
+    public SavedState save_stack(StateMachineBuilder smb, CodeBuilder cob, ClassDesc desc){
+        var name = "LSTATE_"+smb.lstate.size();
+        smb.lstate.add(new StateMachineBuilder.LState(name, desc));
+
+        if(TypeKind.from(desc).slotSize()==2){
+            cob.aload(0).swap().putfield(smb.CD_this, name, desc);
+        }else{
+            cob.aload(0).dup_x2().pop().putfield(smb.CD_this, name, desc);
+        }
+
+        var s = new StackState(name, desc);
+        saved.addLast(s);
+        return s;
+    }
+
+    public SavedState save_local(StateMachineBuilder smb, CodeBuilder cob, ClassDesc desc, int slot){
+        var name = "LSTATE_"+smb.lstate.size();
+        smb.lstate.add(new StateMachineBuilder.LState(name, desc));
+
+        cob.aload(0).loadLocal(TypeKind.from(desc), slot).putfield(smb.CD_this, name, desc);
+
+        var s = new LocalState(name, desc, slot);
+        saved.add(s);
+        return s;
+    }
+
+    public SavedStateTracker restore(StateMachineBuilder smb, CodeBuilder cob, SavedState s){
+        if(!saved.remove(s))throw new IllegalStateException();
+        switch(s){
+            case LocalState(var name, var desc, int slot) ->
+                    cob.aload(0).getfield(smb.CD_this, name, desc).storeLocal(TypeKind.from(desc), slot);
+            case StackState(var name, var desc) ->
+                    cob.aload(0).getfield(smb.CD_this, name, desc);
+        }
+        return this;
+    }
+
+    public void restore_all(StateMachineBuilder smb, CodeBuilder cob) {
+        while(!saved.isEmpty())
+            restore(smb, cob, saved.getFirst());
+    }
+}

src/generator/runtime/SpecialMethodHandler.java

@@ -1,13 +1,26 @@
 package generator.runtime;
 
 import java.lang.classfile.CodeBuilder;
+import java.lang.classfile.Label;
 
-public interface SpecialMethodHandler {
+public abstract class SpecialMethodHandler {
-    void handle(StateMachineBuilder smb, CodeBuilder cob);
+    public final Label handler_start;
+    public final Label handler_resume;
 
-    default boolean removeCall() {
+    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() {
         return true;
     }
 
-    ReplacementKind replacementKind();
+    public void insertShim(StateMachineBuilder smb, CodeBuilder cob){
+        cob.goto_(handler_start).labelBinding(handler_resume);
+    }
+
+    public abstract ReplacementKind replacementKind();
 }

src/generator/runtime/StateMachineBuilder.java

@@ -1,7 +1,5 @@
 package generator.runtime;
 
-import generator.gen.Gen;
-
 import java.lang.classfile.*;
 import java.lang.classfile.attribute.InnerClassInfo;
 import java.lang.classfile.attribute.InnerClassesAttribute;
@@ -30,7 +28,13 @@ public abstract class StateMachineBuilder {
     public final MethodModel src_mem;
     public final CodeModel src_com;
 
-    public LocalTracker lt;
+    ArrayList<LState> lstate = new ArrayList<>();
+
+    private final ArrayList<Frame> frames = new ArrayList<>();
+    private final ArrayList<SpecialMethodHandler> handlers = new ArrayList<>();
+
+    record LState(String name, ClassDesc cd) {
+    }
 
     protected HashMap<SpecialMethod, BiFunction<StateMachineBuilder, CodeBuilder, SpecialMethodHandler>> smmap = new HashMap<>();
 
@@ -64,6 +68,43 @@ public abstract class StateMachineBuilder {
         this.params = mts.parameterArray();
         this.MTD_init = MethodTypeDesc.of(ConstantDescs.CD_void, params);
         this.paramSlotOff = Arrays.stream(params).mapToInt(p -> TypeKind.from(p).slotSize()).sum();
+
+
+        var lt = new FrameTracker(this, src_com);
+        for(var coe : src_com){
+            if(coe instanceof Instruction)
+                frames.add(new Frame(lt.locals(), lt.stack()));
+            lt.encounter(coe);
+        }
+        frames.add(new Frame(lt.locals(), lt.stack()));
+
+        for(var wf : with_frames()){
+            if(wf.coe instanceof Instruction)
+                System.out.println(wf.coe() + " " + wf.frame());
+        }
+    }
+
+    public record WithFrame(CodeElement coe, Frame frame){}
+    public Iterable<WithFrame> with_frames(){
+        return () -> {
+            Iterator<CodeElement> coes = src_com.iterator();
+            return new Iterator<>() {
+                int i = 0;
+
+                @Override
+                public boolean hasNext() {
+                    return coes.hasNext();
+                }
+
+                @Override
+                public WithFrame next() {
+                    var coe = coes.next();
+                    var frame = frames.get(i);
+                    if (coe instanceof Instruction) i++;
+                    return new WithFrame(coe, frame);
+                }
+            };
+        };
     }
 
     public int add_state(Label label) {
@@ -112,6 +153,16 @@ public abstract class StateMachineBuilder {
         });
     }
 
+    public boolean hasAnyHandlers(){
+        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) return true;
+            }
+        }
+        return false;
+    }
+
     protected abstract void buildStateMachineMethod(ClassBuilder clb);
 
     public void buildStateMachineMethodCode(ClassBuilder clb, CodeBuilder cob, int loc_param_off){
@@ -129,33 +180,44 @@ public abstract class StateMachineBuilder {
     public void buildStateMachineCode(ClassBuilder clb, CodeBuilder cob, int loc_param_off) {
         boolean ignore_next_pop = false;
 
-        System.out.println();
+        var invalid_state = cob.newLabel();
-        System.out.println("Begin");
-        this.lt = new LocalTracker(this, src_com, loc_param_off);
-
-        var invalidState = cob.newLabel();
         var start_label = cob.newLabel();
         add_state(start_label);
 
-        var handlers = new ArrayList<SpecialMethodHandler>();
+        for(var wf : with_frames()){
-        for (CodeElement coe : src_com) {
+            if (wf.coe() instanceof InvokeInstruction is){
-            if (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));
             }
         }
-        if(handlers.isEmpty())
+
-            throw new RuntimeException("Not a state machine");
+        cob.aload(0).getfield(CD_this, STATE_NAME, TypeKind.INT.upperBound()).lookupswitch(invalid_state, stateSwitchCases);
-        cob.aload(0).getfield(CD_this, STATE_NAME, TypeKind.INT.upperBound()).lookupswitch(invalidState, stateSwitchCases);
         var start = cob.startLabel();
         var end = cob.newLabel();
         cob.localVariable(0, "this", CD_this, start, end);
 
+        {
+            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);
+                    }
+                }
+            }
+        }
+
         SpecialMethodHandler currentHandler = null;
 
         cob.labelBinding(start_label);
-        for (CodeElement coe : src_com) {
+        for (var wf : with_frames()) {
+            var coe = wf.coe();
+            var frame = wf.frame();
             if (coe instanceof Instruction i) {
                 if (ignore_next_pop)
                     if (i.opcode() == Opcode.POP) {
@@ -164,7 +226,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.handle(this, cob);
+                        currentHandler.insertShim(this, cob);
                         currentHandler = null;
                         continue;
                     }
@@ -175,9 +237,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.handle(this, cob);
+                    if(handler.replacementKind() == ReplacementKind.Immediate) handler.insertShim(this, cob);
                     else if(handler.replacementKind() == ReplacementKind.ImmediateReplacingPop) {
-                        handler.handle(this, cob);
+                        handler.insertShim(this, cob);
                         ignore_next_pop = true;
                     }else
                         currentHandler = handler;
@@ -185,7 +247,6 @@ public abstract class StateMachineBuilder {
                 }
             }
 
-            lt.encounter(coe);
 
 
             switch (coe) {
@@ -204,27 +265,29 @@ public abstract class StateMachineBuilder {
 
                 // convert local function parameters to class fields and offset regular locals
                 case LoadInstruction li when li.slot() < paramSlotOff ->
-                        cob.aload(0).getfield(CD_this, PARAM_PREFIX + li.slot(), lt.paramType(li.slot()));
+                        cob.aload(0).getfield(CD_this, PARAM_PREFIX + li.slot(), frame.locals()[li.slot()].toCD());
                 case LoadInstruction li ->
                         cob.loadLocal(li.typeKind(), li.slot() - paramSlotOff + loc_param_off);
 
                 // convert local function parameters to class fields and offset regular locals
                 case StoreInstruction ls when ls.slot() < paramSlotOff && ls.typeKind().slotSize() == 2 ->
-                        cob.aload(0).dup_x2().pop().putfield(CD_this, PARAM_PREFIX + ls.slot(), lt.paramType(ls.slot()));
+                        cob.aload(0).dup_x2().pop().putfield(CD_this, PARAM_PREFIX + ls.slot(), frame.locals()[ls.slot()].toCD());
                 case StoreInstruction ls when ls.slot() < paramSlotOff ->
-                        cob.aload(0).swap().putfield(CD_this, PARAM_PREFIX + ls.slot(), lt.paramType(ls.slot()));
+                        cob.aload(0).swap().putfield(CD_this, PARAM_PREFIX + ls.slot(), frame.locals()[ls.slot()].toCD());
                 case StoreInstruction ls ->
                         cob.storeLocal(ls.typeKind(), ls.slot() - paramSlotOff + loc_param_off);
 
                 default -> cob.with(coe);
             }
         }
-        cob.labelBinding(invalidState);
+        cob.labelBinding(invalid_state);
         cob.new_(ClassDesc.ofDescriptor(IllegalStateException.class.descriptorString())).dup()
                 .invokespecial(ClassDesc.ofDescriptor(IllegalStateException.class.descriptorString()), ConstantDescs.INIT_NAME, ConstantDescs.MTD_void)
                 .athrow();
         cob.labelBinding(end);
 
-        lt.createLocalStoreFields(clb);
+        for (var lstate : lstate) {
+            clb.withField(lstate.name(), lstate.cd(), ClassFile.ACC_PRIVATE);
+        }
     }
 }

src/generator/runtime/future/FutureSMBuilder.java

@@ -2,11 +2,7 @@ package generator.runtime.future;
 
 import generator.future.Future;
 import generator.future.Waker;
-import generator.gen.Gen;
+import generator.runtime.*;
-import generator.runtime.ReplacementKind;
-import generator.runtime.SpecialMethod;
-import generator.runtime.SpecialMethodHandler;
-import generator.runtime.StateMachineBuilder;
 
 import java.lang.classfile.*;
 import java.lang.constant.ClassDesc;
@@ -26,29 +22,33 @@ public class FutureSMBuilder extends StateMachineBuilder {
 
     public final static String AWAITING_FIELD_NAME = "awaiting";
 
-    static class AwaitHandler implements SpecialMethodHandler{
+    static class AwaitHandler extends SpecialMethodHandler{
-        final int yield_state;
+        final int awaiting_state;
-        final Label yield_label;
+        final Label restore_label;
+        final Label save_label;
 
         public AwaitHandler(StateMachineBuilder smb, CodeBuilder cob) {
-            yield_state = smb.add_state(yield_label = cob.newLabel());
+            super(smb, cob);
+            awaiting_state = smb.add_state(restore_label = cob.newLabel());
+            save_label = cob.newLabel();
         }
 
         public ReplacementKind replacementKind(){return ReplacementKind.Immediate;}
 
         @Override
-        public void handle(StateMachineBuilder smb, CodeBuilder cob) {
+        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
-            System.out.println("Await");
+            cob.aload(0).loadConstant(awaiting_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound());
-            cob.aload(0).loadConstant(yield_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound());
             var start = cob.newBoundLabel();
             cob.dup()
                     .aload(1).invokeinterface(CD_Future, "poll", MTD_Object_Waker).dup()
                     .instanceOf(CD_Pending);
             cob.ifThenElse(bcb -> {
                 bcb.swap().aload(0).swap().putfield(smb.CD_this, AWAITING_FIELD_NAME, CD_Future);
-                smb.lt.savingLocals(smb.CD_this, bcb, () -> {
+
-                    bcb.areturn().labelBinding(yield_label);
+                var saved = frame.save(smb, cob, 2, 0);
-                });
+                bcb.areturn().labelBinding(restore_label);
+                saved.restore_all(smb, cob);
+
                 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);
                 bcb.goto_(start);
@@ -59,37 +59,38 @@ public class FutureSMBuilder extends StateMachineBuilder {
         }
     }
 
-    static class YieldHandler implements SpecialMethodHandler {
+    static class YieldHandler extends SpecialMethodHandler {
         final int resume_state;
         final Label resume_label;
 
         public YieldHandler(StateMachineBuilder smb, CodeBuilder cob) {
+            super(smb, cob);
             resume_state = smb.add_state(resume_label = cob.newLabel());
         }
 
         public ReplacementKind replacementKind(){return ReplacementKind.Immediate;}
 
         @Override
-        public void handle(StateMachineBuilder smb, CodeBuilder cob) {
+        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
-
+            var saved = frame.save(smb, cob, 2, 0);
-            smb.lt.savingLocals(smb.CD_this, cob, () -> {
+            cob.aload(0).loadConstant(resume_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
-                cob.aload(0).loadConstant(resume_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
                         .getstatic(CD_Pending, "INSTANCE", CD_Pending)
                         .areturn();
-                cob.labelBinding(resume_label);
+            saved.restore_all(smb, cob);
-            });
         }
     }
 
-    static class RetHandler implements SpecialMethodHandler{
+    static class RetHandler extends SpecialMethodHandler{
 
-        public RetHandler() {}
+
+        protected RetHandler(StateMachineBuilder smb, CodeBuilder cob) {
+            super(smb, cob);
+        }
 
         public ReplacementKind replacementKind(){return ReplacementKind.ReplacingNextReturn;}
 
         @Override
-        public void handle(StateMachineBuilder smb, CodeBuilder cob) {
+        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
-            System.out.println("Return");
             cob.aload(0).loadConstant(-1).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound()).areturn();
         }
     }
@@ -119,7 +120,7 @@ public class FutureSMBuilder extends StateMachineBuilder {
     public FutureSMBuilder(ClassModel src_clm, MethodModel src_mem, CodeModel src_com) {
         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), (_, _) -> new RetHandler());
+        smmap.put(new SpecialMethod(CD_Future, "ret", MTD_Future_Obj), RetHandler::new);
         smmap.put(new SpecialMethod(CD_Future, "yield", ConstantDescs.MTD_void), YieldHandler::new);
     }
 }

src/generator/runtime/gen/GenSMBuilder.java

@@ -1,10 +1,7 @@
 package generator.runtime.gen;
 
 import generator.gen.Gen;
-import generator.runtime.ReplacementKind;
+import generator.runtime.*;
-import generator.runtime.SpecialMethod;
-import generator.runtime.SpecialMethodHandler;
-import generator.runtime.StateMachineBuilder;
 
 import java.lang.classfile.*;
 import java.lang.constant.ClassDesc;
@@ -23,12 +20,13 @@ 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 implements SpecialMethodHandler {
+    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;
         }
@@ -36,32 +34,33 @@ public class GenSMBuilder extends StateMachineBuilder {
         public ReplacementKind replacementKind(){return ReplacementKind.ImmediateReplacingPop;}
 
         @Override
-        public void handle(StateMachineBuilder smb, CodeBuilder cob) {
+        public void buildHandler(StateMachineBuilder smb, CodeBuilder cob, Frame frame) {
             if(is_void)cob.aconst_null();
 
-            smb.lt.savingLocals(smb.CD_this, cob, () -> {
+//            smb.lt.savingLocals(smb.CD_this, cob, () -> {
-                cob.aload(0).loadConstant(resume_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
+//                cob.aload(0).loadConstant(resume_state).putfield(smb.CD_this, STATE_NAME, TypeKind.INT.upperBound())
-                        .new_(CD_Yield)
+//                        .new_(CD_Yield)
-                        .dup_x1()
+//                        .dup_x1()
-                        .swap()
+//                        .swap()
-                        .invokespecial(CD_Yield, ConstantDescs.INIT_NAME, MethodTypeDesc.of(ConstantDescs.CD_void, ConstantDescs.CD_Object))
+//                        .invokespecial(CD_Yield, ConstantDescs.INIT_NAME, MethodTypeDesc.of(ConstantDescs.CD_void, ConstantDescs.CD_Object))
-                        .areturn();
+//                        .areturn();
-                cob.labelBinding(resume_label);
+//                cob.labelBinding(resume_label);
-            });
+//            });
         }
     }
 
-    static class RetHandler implements SpecialMethodHandler {
+    static class RetHandler extends SpecialMethodHandler {
         final boolean is_void;
 
-        public RetHandler(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 handle(StateMachineBuilder smb, CodeBuilder cob) {
+        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())
@@ -77,8 +76,8 @@ public class GenSMBuilder extends StateMachineBuilder {
         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),(_, _) -> new RetHandler(false));
+        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),(_, _) -> new RetHandler(true));
+        smmap.put(new SpecialMethod(CD_Gen, "ret", MTD_Gen),(smb, cob) -> new RetHandler(smb, cob, true));
     }
 
     @Override