Merge pull request #214 from a16z/ncitron/linker-automation

feat: linker automation

jolt-core/src/benches/bench.rs

@@ -223,7 +223,8 @@ fn prove_example<T: Serialize>(
     input: &T,
 ) -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
     let mut tasks = Vec::new();
-    let mut program = host::Program::new(example_name).input(input);
+    let mut program = host::Program::new(example_name);
+    program.set_input(input);
 
     let task = move || {
         let bytecode = program.decode();

jolt-core/src/host/mod.rs

@@ -1,9 +1,6 @@
 use core::{str::FromStr, u8};
 use std::{
-    collections::HashMap,
-    io::{self, Write},
-    path::PathBuf,
-    process::Command,
+    collections::HashMap, fs::{self, File}, io::{self, Write}, path::PathBuf, process::Command
 };
 
 use ark_ff::PrimeField;
@@ -22,10 +19,13 @@ use crate::jolt::{
     vm::{bytecode::BytecodeRow, rv32i_vm::RV32I},
 };
 
+const DEFAULT_MEMORY_SIZE: usize = 10 * 1024 * 1024;
+
 #[derive(Clone)]
 pub struct Program {
     guest: String,
     input: Vec<u8>,
+    memory_size: usize,
     pub elf: Option<PathBuf>,
 }
 
@@ -34,20 +34,25 @@ impl Program {
         Self {
             guest: guest.to_string(),
             input: Vec::new(),
+            memory_size: DEFAULT_MEMORY_SIZE,
             elf: None,
         }
     }
 
-    pub fn input<T: Serialize>(mut self, input: &T) -> Self {
+    pub fn set_input<T: Serialize>(&mut self, input: &T) {
         let mut serialized = postcard::to_stdvec(input).unwrap();
         self.input.append(&mut serialized);
+    }
 
-        self
+    pub fn set_memory_size(&mut self, len: usize) {
+        self.memory_size = len;
     }
 
     pub fn build(&mut self) {
         if self.elf.is_none() {
+            self.save_linker(self.memory_size);
             let output = Command::new("cargo")
+                .envs([("RUSTFLAGS", format!("-C link-arg=-T{}", self.linker_path()))])
                 .args(&[
                     "build",
                     "--release",
@@ -158,4 +163,49 @@ impl Program {
 
         device.outputs
     }
+
+    fn save_linker(&self, memory_size: usize) {
+        let linker_path = PathBuf::from_str(&self.linker_path()).unwrap();
+        if let Some(parent) = linker_path.parent() {
+            fs::create_dir_all(parent).expect("could not create linker file");
+        }
+
+        let linker_script = LINKER_SCRIPT_TEMPLATE
+            .replace("{MEMORY_SIZE}", &memory_size.to_string());
+
+        let mut file = File::create(linker_path).expect("could not create linker file");
+        file.write(linker_script.as_bytes()).expect("could not save linker");
+    }
+
+    fn linker_path(&self) -> String {
+        format!("/tmp/jolt-guest-linkers/{}.ld", self.guest)
+    }
+}
+
+const LINKER_SCRIPT_TEMPLATE: &str = r#"
+MEMORY {
+  program (rwx) : ORIGIN = 0x80000000, LENGTH = {MEMORY_SIZE}
+}
+
+SECTIONS {
+  .text.boot : {
+    *(.text.boot)
+  } > program
+
+  .text : {
+    *(.text)
+  } > program
+
+  .data : {
+    *(.data)
+  } > program
+
+  .bss : {
+    *(.bss)
+  } > program
+
+  . = ALIGN(8);
+  . = . + 4096;
+  _STACK_PTR = .;
 }
+"#;

jolt-core/src/jolt/vm/rv32i_vm.rs

@@ -226,7 +226,8 @@ mod tests {
     fn fib_e2e() {
         let _guard = FIB_FILE_LOCK.lock().unwrap();
 
-        let mut program = host::Program::new("fibonacci-guest").input(&9u32);
+        let mut program = host::Program::new("fibonacci-guest");
+        program.set_input(&9u32);
         let bytecode = program.decode();
         let (io_device, bytecode_trace, instruction_trace, memory_trace, circuit_flags) =
             program.trace();
@@ -253,7 +254,8 @@ mod tests {
     fn sha3_e2e() {
         let _guard = SHA3_FILE_LOCK.lock().unwrap();
 
-        let mut program = host::Program::new("sha3-guest").input(&[5u8; 32]);
+        let mut program = host::Program::new("sha3-guest");
+        program.set_input(&[5u8; 32]);
         let bytecode = program.decode();
         let (io_device, bytecode_trace, instruction_trace, memory_trace, circuit_flags) =
             program.trace();

jolt-sdk/macros/Cargo.toml

@@ -18,6 +18,7 @@ proc-macro = true
 [dependencies]
 syn = { version = "1.0", features = ["full"] }
 quote = "1.0"
+proc-macro2 = "1.0.79"
 
 postcard = "1.0.8"
 serde = { version = "1.0.196", features = [] }