Removing hook between encode and export

2 files changed, 80 insertions, 19 deletions

diff --git a/chapter28b/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java b/chapter28b/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java
index 81d99190..67176710 100644
--- a/chapter28b/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java
+++ b/chapter28b/src/main/java/com/seaofnodes/simple/codegen/CodeGen.java
@@ -227,7 +227,8 @@ public class CodeGen {
 
     // Run requested phases. File export must request Serialize before Lower;
     // requesting Export up front does that automatically. Encoding alone does
-    // not retain reloadable ideal IR.
+    // not retain reloadable ideal IR. Choose combined in-memory output before
+    // Encoding; otherwise saved IR requests separate object encodings.
 
     // No code emission, just IR generation
     public CodeGen driver( Phase phase ) { return driver(phase,null,null,false,false,0); }
@@ -247,7 +248,6 @@ public class CodeGen {
 
     // The Generic driver
     public CodeGen driver( Phase target, String cpu, String callingConv, boolean inMemory, boolean emitEntrySymbol, int dump ) {
-        _splitObjects = target.ordinal()>=Phase.Export.ordinal() && !inMemory;
         // Parse a Simple file, and for mentioned external symbols either load
         // from .o or recursively parse
         step(target, Phase.Parse,        this::parse,        dump);
@@ -281,7 +281,7 @@ public class CodeGen {
         // Register allocation
         step(target, Phase.RegAlloc,     this::regAlloc,     dump);
         // Encode machine instructions into bits
-        step(target, Phase.Encoding,     this::encode,       dump);
+        step(target, Phase.Encoding, () -> encode(inMemory), dump);
         // Export bits (and serialized IR) into ELF files
         step(target, Phase.Export, () -> exportELF(inMemory,emitEntrySymbol), dump);
         return this;
@@ -456,7 +456,6 @@ public class CodeGen {
     CompUnit _emitUnit;
     boolean emits(CompUnit cu) { return cu!=null && cu._src!=null && (_emitUnit==null || _emitUnit==cu); }
     boolean emits(FunNode fun) { return emits(fun._compunit); }
-    private boolean _splitObjects;
     final LinkedHashMap<CompUnit,Encoding> _encodings = new LinkedHashMap<>();
 
 
@@ -674,16 +673,19 @@ public class CodeGen {
     // ---------------------------
     // Encoding
     public Encoding _encoding;
-    public CodeGen encode() { return phase(Phase.Encoding, _phase==Phase.RegAlloc, () -> {
-        if( !_splitObjects ) { _encoding = new Encoding(this).encode(); return; }
-        try {
-            for( CompUnit cu : _serials.keySet() ) {
-                _emitUnit = cu;
-                Encoding enc = new Encoding(this).encode();
-                _encodings.put(cu,enc);
-                if( cu._fname.equals(_srcName) ) _encoding = enc;
-            }
-        } finally { _emitUnit = null; }
+    public CodeGen encode() { return encode(false); }
+    private CodeGen encode(boolean inMemory) { return phase(Phase.Encoding, _phase==Phase.RegAlloc, () -> {
+        // Saved ideal IR requests one object per unit, including when the
+        // driver stopped at Serialize and resumed later. In-memory linking
+        // explicitly requests one combined image instead.
+        if( inMemory || _serials.isEmpty() ) { _encoding = new Encoding(this).encode(); return; }
+        for( CompUnit cu : _serials.keySet() ) {
+            _emitUnit = cu;
+            Encoding enc = new Encoding(this).encode();
+            _encodings.put(cu,enc);
+            if( cu._fname.equals(_srcName==null ? "Test" : _srcName) ) _encoding = enc;
+        }
+        _emitUnit = null;
     }); }
 
     // ---------------------------
@@ -691,17 +693,17 @@ public class CodeGen {
     ElfWriter _elf;
     public CodeGen exportELF( boolean inMemory, boolean emitEntrySymbol ) {
         return phase(Phase.Export, _phase==Phase.Encoding, () -> {
-            if( inMemory ) new LinkMem(this).link(_encoding);
+            if( inMemory ) {
+                assert _encodings.isEmpty() : "In-memory output must be selected before Encoding";
+                new LinkMem(this).link(_encoding);
+            }
             else {
                 if( _serials.isEmpty() )
                     throw new IllegalStateException("ELF export needs ideal IR: request Export from the driver, or Serialize before Lower");
+                assert !_encodings.isEmpty() : "File output must be selected before Encoding";
                 Encoding entry = _encoding;
                 BAOS serial = _serial;
                 try {
-                    if( _encodings.isEmpty() ) for( CompUnit cu : _serials.keySet() ) {
-                        _emitUnit=cu;
-                        _encodings.put(cu,new Encoding(this).encode());
-                    }
                     for( var e : _encodings.entrySet() ) {
                         _emitUnit = e.getKey();
                         _encoding = e.getValue();
diff --git a/chapter28b/src/test/java/com/seaofnodes/simple/PhaseTest.java b/chapter28b/src/test/java/com/seaofnodes/simple/PhaseTest.java
index f2d09c7c..6505009d 100644
--- a/chapter28b/src/test/java/com/seaofnodes/simple/PhaseTest.java
+++ b/chapter28b/src/test/java/com/seaofnodes/simple/PhaseTest.java
@@ -4,6 +4,10 @@ import com.seaofnodes.graph.GraphObserver;
 import com.seaofnodes.simple.codegen.CodeGen;
 import com.seaofnodes.simple.codegen.CodeGen.Phase;
 import com.seaofnodes.simple.node.Node;
+import com.seaofnodes.simple.node.FunNode;
+import com.seaofnodes.simple.type.TypeInteger;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import org.junit.Test;
 import static org.junit.Assert.*;
@@ -71,11 +75,19 @@ public class PhaseTest {
         assertEquals(Phase.Serialize,code._phase);
         code.driver(Phase.Encoding,"riscv","SystemV");
         assertNotNull(code._serial);
+        assertNotNull(code._encoding);
 
         CodeGen memory = new CodeGen(SRC).driver("riscv","SystemV",true,false);
         assertNull(memory._serial);
         assertEquals(Phase.Export,memory._phase);
 
+        // Keeping ideal IR must not force separate images when the caller
+        // explicitly requests in-memory linking.
+        CodeGen saved = new CodeGen(SRC).driver(Phase.Serialize);
+        saved.driver("riscv","SystemV",true,false);
+        assertNotNull(saved._serial);
+        assertArrayEquals(memory._encoding._bits.toByteArray(),saved._encoding._bits.toByteArray());
+
         CodeGen late = new CodeGen(SRC).driver(Phase.Encoding,"riscv","SystemV");
         try {
             late.exportELF(false,false);
@@ -84,4 +96,51 @@ public class PhaseTest {
             assertTrue(expected.getMessage().contains("Serialize before Lower"));
         }
     }
+
+    @Test public void testStagedObjectExport() throws Exception {
+        Path root = Files.createTempDirectory(Files.createDirectories(Path.of("build/objs")),"phase-export-");
+        Path src = Files.createDirectories(root.resolve("src"));
+        Files.createDirectories(src.resolve("Main"));
+        Files.writeString(src.resolve("Main.smp"),"val run_noInline = { int n -> Helper.twice_noInline(n)+1; };");
+        Files.writeString(src.resolve("Main/Helper.smp"),"val twice_noInline = { int n -> n*2; };");
+        Path cfile = root.resolve("driver.c");
+        Files.writeString(cfile,"""
+            __attribute__((CALL_CONV)) extern long long run_noInline(long long);
+            int main(void) { return run_noInline(20)!=41 || run_noInline(-3)!=-5; }
+            """);
+        for( boolean direct : new boolean[]{false,true} ) {
+            Path obj = Files.createDirectories(root.resolve(direct ? "direct" : "driver"));
+            CodeGen code = new CodeGen(src.toString(),obj.toString(),null,"Main",null,123,true,TypeInteger.BOT)
+                .driver(Phase.Serialize);
+            if( direct ) {
+                code.driver(Phase.RegAlloc,TestC.CPU_PORT,TestC.CALL_CONVENTION);
+                code.encode();
+            } else {
+                code.driver(Phase.Encoding,TestC.CPU_PORT,TestC.CALL_CONVENTION);
+            }
+            // Encoding has already separated the entry object's code from
+            // its dependency. Export only writes the prepared encodings.
+            assertNotNull(code._encoding);
+            for( var block : code._encoding._cfg )
+                if( block instanceof FunNode fun )
+                    assertSame(code._compunits.get("Main"),fun._compunit);
+            assertFalse(Files.exists(obj.resolve("Main.o")));
+            if( direct ) code.exportELF(false,false);
+            else code.driver(TestC.CPU_PORT,TestC.CALL_CONVENTION,false,false);
+            assertTrue(Files.exists(obj.resolve("Main.o")));
+            assertTrue(Files.exists(obj.resolve("Main/Helper.o")));
+            String exe = obj.resolve(TestC.OS.startsWith("Windows") ? "run.exe" : "run").toString();
+            assertEquals("",TestC.gcc(code,TestC.CALL_CONVENTION.equals("win64") ? "ms_abi" : "sysv_abi",
+                                     cfile.toString(),null,exe));
+        }
+        // The same serialized units can instead form one linked memory image.
+        CodeGen memory = new CodeGen(src.toString(),root.resolve("memory").toString(),null,
+            "Main",null,123,true,TypeInteger.BOT).driver(Phase.Serialize);
+        memory.driver("riscv","SystemV",true,false);
+        boolean helper = false;
+        for( var block : memory._encoding._cfg )
+            if( block instanceof FunNode fun && fun._compunit==memory._compunits.get("Main/Helper") )
+                helper = true;
+        assertTrue("The combined image includes dependency bodies",helper);
+    }
 }