✨ Code tables

README.md

@@ -13,8 +13,6 @@ in Rust. It comes with:
 - [ ] Jump tables
     - [ ] parsing
     - [ ] builtins (`__tablestart`, `__tablesize`)
-- [ ] Code tables
-    - [ ] builtins (`__tablestart`, `__tablesize`)
 - [ ] ABI builtins (`__FUNC_SIG`, `__EVEN_HASH`, `__ERROR`)
 - [ ] Imports (`#include` statements)
 

crates/analysis/src/errors.rs

@@ -63,8 +63,8 @@ pub enum AnalysisError<'ast, 'src> {
     },
 }
 
-impl AnalysisError<'_, '_> {
-    pub fn report(&self, filename: String) -> Report<(String, std::ops::Range<usize>)> {
+impl<'ast, 'src> AnalysisError<'ast, 'src> {
+    pub fn report(&self, filename: String) -> Report<'_, (String, std::ops::Range<usize>)> {
         match self {
             Self::DefinitionNameCollision {
                 collided,

crates/analysis/src/lib.rs

@@ -249,7 +249,7 @@ impl<'a, 'src, 'ast: 'src, E: FnMut(AnalysisError<'ast, 'src>)> MacroAnalysis<'a
                     if !global_exists!(
                         self.global_defs,
                         table_ref.ident(),
-                        Definition::Table { .. } | Definition::Jumptable(_)
+                        Definition::CodeTable { .. } | Definition::Jumptable(_)
                     ) {
                         self.emit(AnalysisError::DefinitionNotFound {
                             scope: self.m,
@@ -258,10 +258,16 @@ impl<'a, 'src, 'ast: 'src, E: FnMut(AnalysisError<'ast, 'src>)> MacroAnalysis<'a
                         })
                     }
 
-                    self.emit(AnalysisError::NotYetSupported {
-                        intent: "__tablesize and __tableoffset".to_string(),
-                        span: ((), table_ref.1),
-                    });
+                    if global_exists!(
+                        self.global_defs,
+                        table_ref.ident(),
+                        Definition::Jumptable(_)
+                    ) {
+                        self.emit(AnalysisError::NotYetSupported {
+                            intent: "__tablesize and __tableoffset on jumptable".to_string(),
+                            span: ((), table_ref.1),
+                        });
+                    }
                 }
                 Invoke::BuiltinCodeSize(code_ref) | Invoke::BuiltinCodeOffset(code_ref) => {
                     if !global_exists!(self.global_defs, code_ref.ident(), Definition::Macro(_)) {
@@ -473,7 +479,7 @@ mod test {
             expr: (ConstExpr::FreeStoragePointer, span),
         };
 
-        let unrelated_table = Definition::Table {
+        let unrelated_table = Definition::CodeTable {
             name: ("awesome_stuff", span),
             data: Box::new([0x00, 0x01]),
         };

crates/ast/src/ast.rs

@@ -20,7 +20,7 @@ pub enum Definition<'src> {
         expr: Spanned<ConstExpr>,
     },
     Jumptable(Jumptable<'src>),
-    Table {
+    CodeTable {
         name: Spanned<&'src str>,
         data: Box<[u8]>,
     },
@@ -47,7 +47,7 @@ impl<'src> IdentifiableNode<'src> for Definition<'src> {
             Self::Macro(m) => &m.name,
             Self::Constant { name, .. } => name,
             Self::Jumptable(jt) => &jt.name,
-            Self::Table { name, .. } => name,
+            Self::CodeTable { name, .. } => name,
             Self::SolEvent(e) => &e.name,
             Self::SolError(e) => &e.name,
             Self::SolFunction(f) => &f.name,

crates/ast/src/lexer.rs

@@ -9,7 +9,7 @@ use chumsky::{
 use std::fmt;
 
 /// Lex the given source code string into tokens.
-pub(crate) fn lex(src: &str) -> Result<Vec<Spanned<Token>>, Vec<Rich<'_, Token<'_>>>> {
+pub(crate) fn lex<'a>(src: &'a str) -> Result<Vec<Spanned<Token<'a>>>, Vec<Rich<'a, Token<'a>>>> {
     lexer().parse(src).into_result().map_err(|e| {
         e.into_iter()
             .map(|errs| errs.map_token(Token::Error))

crates/ast/src/parser.rs

@@ -128,7 +128,7 @@ fn macro_statement<'tokens, 'src: 'tokens>() -> impl Parser<'tokens, 'src, ast::
 }
 
 fn instruction<'tokens, 'src: 'tokens>() -> impl Parser<'tokens, 'src, ast::Instruction<'src>> {
-    let push_auto = word().map(|(value, span)| (ast::Instruction::VariablePush((value, span))));
+    let push_auto = word().map(|(value, span)| ast::Instruction::VariablePush((value, span)));
 
     let push = select! {
         Ident("push1") => 1,
@@ -280,7 +280,7 @@ fn table<'tokens, 'src: 'tokens>() -> impl Parser<'tokens, 'src, ast::Definition
                 .collect::<Vec<_>>()
                 .delimited_by(punct('{'), punct('}')),
         )
-        .map(|(name, code)| ast::Definition::Table {
+        .map(|(name, code)| ast::Definition::CodeTable {
             name,
             data: code
                 .into_iter()
@@ -668,7 +668,7 @@ mod tests {
         assert_ok!(
             table(),
             vec![Ident("table"), Ident("TEST"), Punct('{'), Hex("0xc0de"), Punct('}')],
-            ast::Definition::Table {
+            ast::Definition::CodeTable {
                 name: ("TEST", span),
                 data: Box::new([0xc0, 0xde])
             }
@@ -683,7 +683,7 @@ mod tests {
                 Hex("0xcc00ddee"),
                 Punct('}')
             ],
-            ast::Definition::Table {
+            ast::Definition::CodeTable {
                 name: ("TEST", span),
                 data: Box::new([0xc0, 0xde, 0xcc, 0x00, 0xdd, 0xee])
             }

crates/compilation/src/lib.rs

@@ -27,14 +27,46 @@ impl IncludedMacro<'_> {
     fn start_ref(&self) -> MarkRef {
         MarkRef {
             ref_type: RefType::Direct(self.start_id),
+
             is_pushed: true,
             set_size: None,
         }
     }
 }
 
-pub fn generate_for_entrypoint<'src>(
-    globals: &mut CompileGlobals<'src, '_>,
+struct IncludedCodeTable<'src, 'ast> {
+    name: &'src str,
+    referenced: bool,
+    start_id: usize,
+    end_id: usize,
+    data: &'ast [u8],
+}
+
+struct ProgramDataDeps<'src, 'ast> {
+    included_macros: Vec<IncludedMacro<'src>>,
+    included_code_tables: Vec<IncludedCodeTable<'src, 'ast>>,
+}
+
+impl<'src, 'ast> IncludedCodeTable<'src, 'ast> {
+    fn size_ref(&self) -> MarkRef {
+        MarkRef {
+            ref_type: RefType::Delta(self.start_id, self.end_id),
+            is_pushed: true,
+            set_size: None,
+        }
+    }
+
+    fn start_ref(&self) -> MarkRef {
+        MarkRef {
+            ref_type: RefType::Direct(self.start_id),
+            is_pushed: true,
+            set_size: None,
+        }
+    }
+}
+
+pub fn generate_for_entrypoint<'src, 'ast: 'src>(
+    globals: &mut CompileGlobals<'src, 'ast>,
     entry_point: &Macro<'src>,
 ) -> Vec<u8> {
     let mut mark_tracker = MarkTracker::default();
@@ -48,6 +80,27 @@ pub fn generate_for_entrypoint<'src>(
         start_id,
         end_id,
     });
+    let included_code_tables: Vec<IncludedCodeTable<'src, 'ast>> = globals
+        .defs
+        .iter()
+        .filter_map(|(_, def)| {
+            let Definition::CodeTable { name, data } = def else {
+                return None;
+            };
+            Some(IncludedCodeTable {
+                name: name.ident(),
+                referenced: false,
+                data,
+                start_id: mark_tracker.next_mark(),
+                end_id: mark_tracker.next_mark(),
+            })
+        })
+        .collect();
+
+    let mut program_data_deps = ProgramDataDeps {
+        included_macros,
+        included_code_tables,
+    };
 
     let mut asm = Vec::with_capacity(10_000);
     asm.push(Asm::Mark(start_id));
@@ -57,21 +110,35 @@ pub fn generate_for_entrypoint<'src>(
         Box::new([]),
         &mut mark_tracker,
         &mut label_stack,
-        &mut included_macros,
+        &mut program_data_deps,
         &mut asm,
     );
 
-    included_macros.into_iter().skip(1).for_each(|included| {
-        let section_macro =
-            if let Some(Definition::Macro(section_macro)) = globals.defs.get(included.name) {
-                section_macro
-            } else {
-                panic!("Section macro {} not found", included.name);
-            };
-        asm.push(Asm::Mark(included.start_id));
-        asm.push(Asm::Data(generate_for_entrypoint(globals, section_macro)));
-        asm.push(Asm::Mark(included.end_id));
-    });
+    program_data_deps
+        .included_macros
+        .into_iter()
+        .skip(1)
+        .for_each(|included| {
+            let section_macro =
+                if let Some(Definition::Macro(section_macro)) = globals.defs.get(included.name) {
+                    section_macro
+                } else {
+                    panic!("Section macro {} not found", included.name);
+                };
+            asm.push(Asm::Mark(included.start_id));
+            asm.push(Asm::Data(generate_for_entrypoint(globals, section_macro)));
+            asm.push(Asm::Mark(included.end_id));
+        });
+
+    program_data_deps
+        .included_code_tables
+        .into_iter()
+        .filter(|t| t.referenced)
+        .for_each(|included| {
+            asm.push(Asm::Mark(included.start_id));
+            asm.push(Asm::Data(included.data.to_vec()));
+            asm.push(Asm::Mark(included.end_id));
+        });
 
     asm.push(Asm::Mark(end_id));
 
@@ -132,13 +199,13 @@ pub fn generate_default_constructor(runtime: Vec<u8>) -> Box<[Asm]> {
     .into_boxed_slice()
 }
 
-fn generate_for_macro<'src: 'cmp, 'cmp>(
-    globals: &mut CompileGlobals<'src, '_>,
+fn generate_for_macro<'src: 'cmp, 'cmp, 'ast>(
+    globals: &mut CompileGlobals<'src, 'ast>,
     current: &Macro<'src>,
     arg_values: Box<[Asm]>,
     mark_tracker: &mut MarkTracker,
     label_stack: &'cmp mut LabelStack<'src, usize>,
-    included_macros: &'cmp mut Vec<IncludedMacro<'src>>,
+    program_data_deps: &'cmp mut ProgramDataDeps<'src, 'ast>,
     asm: &mut Vec<Asm>,
 ) {
     let current_args: BTreeMap<&str, Asm> = BTreeMap::from_iter(
@@ -182,13 +249,15 @@ fn generate_for_macro<'src: 'cmp, 'cmp>(
                         .collect(),
                     mark_tracker,
                     label_stack,
-                    included_macros,
+                    program_data_deps,
                     asm,
                 )
             }
             Invoke::BuiltinCodeSize(code_ref) => {
-                let mref: MarkRef = if let Some(included) =
-                    included_macros.iter().find(|m| m.name == code_ref.ident())
+                let mref: MarkRef = if let Some(included) = program_data_deps
+                    .included_macros
+                    .iter()
+                    .find(|m| m.name == code_ref.ident())
                 {
                     included.size_ref()
                 } else {
@@ -200,14 +269,16 @@ fn generate_for_macro<'src: 'cmp, 'cmp>(
                         end_id,
                     };
                     let mref = included.size_ref();
-                    included_macros.push(included);
+                    program_data_deps.included_macros.push(included);
                     mref
                 };
                 asm.push(Asm::Ref(mref));
             }
             Invoke::BuiltinCodeOffset(code_ref) => {
-                let mref: MarkRef = if let Some(included) =
-                    included_macros.iter().find(|m| m.name == code_ref.ident())
+                let mref: MarkRef = if let Some(included) = program_data_deps
+                    .included_macros
+                    .iter()
+                    .find(|m| m.name == code_ref.ident())
                 {
                     included.start_ref()
                 } else {
@@ -219,11 +290,29 @@ fn generate_for_macro<'src: 'cmp, 'cmp>(
                         end_id,
                     };
                     let mref = included.start_ref();
-                    included_macros.push(included);
+                    program_data_deps.included_macros.push(included);
                     mref
                 };
                 asm.push(Asm::Ref(mref));
             }
+            Invoke::BuiltinTableStart(table_ref) => {
+                let target_table = program_data_deps
+                    .included_code_tables
+                    .iter_mut()
+                    .find(|t| t.name == table_ref.ident())
+                    .expect("Table not found (might be jumptable)");
+                target_table.referenced = true;
+                asm.push(Asm::Ref(target_table.start_ref()));
+            }
+            Invoke::BuiltinTableSize(table_ref) => {
+                let target_table = program_data_deps
+                    .included_code_tables
+                    .iter_mut()
+                    .find(|t| t.name == table_ref.ident())
+                    .expect("Table not found (might be jumptable)");
+                target_table.referenced = true;
+                asm.push(Asm::Ref(target_table.size_ref()));
+            }
             _ => panic!(
                 "Compilation not yet implemented for this invocation type `{:?}`",
                 invoke

examples/Ownable.huff

@@ -0,0 +1,27 @@
+#define constant OWNABLE_SLOT = FREE_STORAGE_POINTER()
+
+#define macro CONSTRUCTOR() = {
+    caller [OWNABLE_SLOT] sstore
+    _RETURN_MAIN()
+}
+
+#define macro _RETURN_MAIN() = {
+    __codesize(MAIN)        // [size]
+    dup1                    // [size, size]
+    __codeoffset(MAIN)      // [offset, size, size]
+    0x0                     // [0, offset, size, size]
+    codecopy                // [size]
+    0x0                     // [0, size]
+    return                  // []
+}
+
+#define macro MAIN() = {
+    [OWNABLE_SLOT] sload    // [owner]
+    caller                  // [msg.sender, owner]
+    eq                      // [msg.sender == owner]
+    is_owner jumpi          // []
+        0x0 0x0 revert
+    is_owner:
+
+    stop
+}

examples/features/CodeTable.huff

@@ -0,0 +1,28 @@
+#define macro CONSTRUCTOR() = {
+    0x1 0x0 sstore
+    __codesize(MAIN)
+    dup1
+    __codeoffset(MAIN)
+    0x0
+    codecopy
+    0x0
+    return
+}
+
+
+
+
+#define macro MAIN() = takes(0) returns(0) {
+    __tablestart(CODE)
+    __tablestart(WOW)
+
+    __tablesize(CODE)
+}
+
+#define table CODE {
+    0xaa11bb0000000000000000000000000000000000000000000000000000000000000000
+}
+
+#define table WOW {
+    0x123456
+}