Repository: huff-language/huff-rs
Pull request #171 (2022-07-21 14:19:38)

feat(huff_cli): Constant Overrides
## Overview

Adds the ability to override / set new constant definitions for the current compilation environment via the `-c` flag in the CLI.

#### Example:
`huffc ./Test.huff -b -c a=0x01 b=0xa57b`



Patch:
diff --git a/huff_cli/src/huffc.rs b/huff_cli/src/huffc.rs
index c20a643..cf9a596 100644
--- a/huff_cli/src/huffc.rs
+++ b/huff_cli/src/huffc.rs
@@ -12,12 +12,12 @@ use ethers_core::utils::hex;
 use huff_codegen::Codegen;
 use huff_core::Compiler;
 use huff_utils::prelude::{
-    export_interfaces, gen_sol_interfaces, unpack_files, AstSpan, CodegenError, CodegenErrorKind,
-    CompilerError, FileSource, OutputLocation, Span,
+    export_interfaces, gen_sol_interfaces, str_to_bytes32, unpack_files, AstSpan, CodegenError,
+    CodegenErrorKind, CompilerError, FileSource, Literal, OutputLocation, Span,
 };
 use isatty::stdout_isatty;
 use spinners::{Spinner, Spinners};
-use std::{io::Write, path::Path, sync::Arc};
+use std::{collections::BTreeMap, io::Write, path::Path, sync::Arc};
 use yansi::Paint;
 
 /// The Huff CLI Args
@@ -70,6 +70,10 @@ struct Huff {
     /// Verbose output.
     #[clap(short = 'v', long = "verbose")]
     verbose: bool,
+
+    /// Override / set constants for the compilation environment.
+    #[clap(short = 'c', long = "constants", multiple_values = true)]
+    constants: Option<Vec<String>>,
 }
 
 /// Helper function to read an stdin input
@@ -102,6 +106,33 @@ fn main() {
         }
     };
 
+    // If constant overrides were passed, create a map of their names and values
+    let constants: Option<BTreeMap<&str, Literal>> = cli.constants.as_ref().map(|_constants| {
+        _constants
+            .iter()
+            .map(|c: &String| {
+                let parts = c.as_str().split('=').collect::<Vec<_>>();
+
+                // Check that constant override argument is valid
+                // Key rule: Alphabetic chars + underscore
+                // Value rule: Valid literal string (0x...)
+                if parts.len() != 2 ||
+                    parts[0].chars().any(|c| !(c.is_alphabetic() || c == '_')) ||
+                    !parts[1].starts_with("0x") ||
+                    parts[1][2..].chars().any(|c| {
+                        !(c.is_numeric() ||
+                            matches!(c, '\u{0041}'..='\u{0046}' | '\u{0061}'..='\u{0066}'))
+                    })
+                {
+                    eprintln!("Invalid constant override argument: {}", Paint::red(c.to_string()));
+                    std::process::exit(1);
+                }
+
+                (parts[0], str_to_bytes32(&parts[1][2..]))
+            })
+            .collect()
+    });
+
     let mut use_cache = true;
     if cli.interactive {
         // Don't accept configured inputs
@@ -124,6 +155,7 @@ fn main() {
         sources: Arc::clone(&sources),
         output,
         construct_args: cli.inputs,
+        constant_overrides: constants,
         optimize: cli.optimize,
         bytecode: cli.bytecode,
         cached: use_cache,
diff --git a/huff_codegen/README.md b/huff_codegen/README.md
index a46ec3c..4a63122 100644
--- a/huff_codegen/README.md
+++ b/huff_codegen/README.md
@@ -27,6 +27,8 @@ Below we showcase generating a compile artifact from compiled bytecode using `hu
 use huff_codegen::*;
 use huff_utils::files::FileSource;
 use std::sync::Arc;
+use std::cell::RefCell;
+use std::rc::Rc;
 
 // Instantiate an empty Codegen
 let mut cg = Codegen::new();
@@ -53,6 +55,8 @@ Let's say you have a [Contract](../huff_utils/ast/struct.Contract.html) instance
 use huff_codegen::*;
 use huff_utils::prelude::*;
 use std::sync::Arc;
+use std::cell::RefCell;
+use std::rc::Rc;
 
 // Mock contract with a main macro
 let contract = Contract {
@@ -86,7 +90,7 @@ let contract = Contract {
   ],
   invocations: vec![],
   imports: vec![],
-  constants: vec![],
+  constants: Rc::new(RefCell::new(vec![])),
   functions: vec![],
   events: vec![],
   tables: vec![],
@@ -105,6 +109,8 @@ Similarly, once you have a [Contract](../huff_utils/ast/struct.Contract.html) in
 use huff_codegen::*;
 use huff_utils::prelude::*;
 use std::sync::Arc;
+use std::cell::RefCell;
+use std::rc::Rc;
 
 // Mock contract with a constructor macro
 let contract = Contract {
@@ -138,7 +144,7 @@ let contract = Contract {
   ],
   invocations: vec![],
   imports: vec![],
-  constants: vec![],
+  constants: Rc::new(RefCell::new(vec![])),
   functions: vec![],
   events: vec![],
   tables: vec![],
diff --git a/huff_codegen/src/irgen/arg_calls.rs b/huff_codegen/src/irgen/arg_calls.rs
index 6e62d55..f1bc6f5 100644
--- a/huff_codegen/src/irgen/arg_calls.rs
+++ b/huff_codegen/src/irgen/arg_calls.rs
@@ -22,7 +22,8 @@ pub fn bubble_arg_call(
     let starting_offset = *offset;
 
     // Check Constant Definitions
-    if let Some(constant) = contract.constants.iter().find(|const_def| const_def.name.eq(arg_name))
+    if let Some(constant) =
+        contract.constants.borrow().iter().find(|const_def| const_def.name.eq(arg_name))
     {
         tracing::info!(target: "codegen", "ARGCALL IS CONSTANT: {:?}", constant);
         let push_bytes = match &constant.value {
diff --git a/huff_codegen/src/irgen/constants.rs b/huff_codegen/src/irgen/constants.rs
index a4a68c4..ab0bfc3 100644
--- a/huff_codegen/src/irgen/constants.rs
+++ b/huff_codegen/src/irgen/constants.rs
@@ -9,18 +9,18 @@ pub fn constant_gen(
     ir_byte_span: AstSpan,
 ) -> Result<String, CodegenError> {
     // Get the first `ConstantDefinition` that matches the constant's name
-    let constant =
-        if let Some(m) = contract.constants.iter().find(|const_def| const_def.name.eq(&name)) {
-            m
-        } else {
-            tracing::error!(target: "codegen", "MISSING CONSTANT DEFINITION \"{}\"", name);
+    let constants = contract.constants.borrow();
+    let constant = if let Some(m) = constants.iter().find(|const_def| const_def.name.eq(&name)) {
+        m
+    } else {
+        tracing::error!(target: "codegen", "MISSING CONSTANT DEFINITION \"{}\"", name);
 
-            return Err(CodegenError {
-                kind: CodegenErrorKind::MissingConstantDefinition(name.to_string()),
-                span: ir_byte_span,
-                token: None,
-            })
-        };
+        return Err(CodegenError {
+            kind: CodegenErrorKind::MissingConstantDefinition(name.to_string()),
+            span: ir_byte_span,
+            token: None,
+        })
+    };
 
     // Generate bytecode for the constant
     // Should always be a `Literal` if storage pointers were derived in the AST
diff --git a/huff_core/README.md b/huff_core/README.md
index 880cbb3..de17c6a 100644
--- a/huff_core/README.md
+++ b/huff_core/README.md
@@ -19,9 +19,11 @@ use huff_core::Compiler;
 use huff_utils::error::CompilerError;
 use huff_utils::artifact::Artifact;
 use std::sync::Arc;
+use std::cell::RefCell;
+use std::rc::Rc;
 
 // Instantiate the Compiler Instance
-let mut compiler = Compiler::new(Arc::new(vec!["../huff-examples/erc20/contracts/ERC20.huff".to_string()]), None, None, false, false);
+let mut compiler = Compiler::new(Arc::new(vec!["../huff-examples/erc20/contracts/ERC20.huff".to_string()]), None, None, None, false, false);
 
 // Execute the compiler
 let res: Result<Vec<Arc<Artifact>>, Arc<CompilerError<'_>>> = compiler.execute();
diff --git a/huff_core/src/lib.rs b/huff_core/src/lib.rs
index 4003086..4e595ae 100644
--- a/huff_core/src/lib.rs
+++ b/huff_core/src/lib.rs
@@ -11,6 +11,7 @@ use huff_parser::*;
 use huff_utils::prelude::*;
 use rayon::prelude::*;
 use std::{
+    collections::BTreeMap,
     ffi::OsString,
     fs,
     path::{Path, PathBuf},
@@ -42,18 +43,21 @@ pub(crate) mod cache;
 ///     Arc::new(vec!["../huff-examples/erc20/contracts/ERC20.huff".to_string()]),
 ///     Some("./artifacts".to_string()),
 ///     None,
+///     None,
 ///     false,
 ///     false
 /// );
 /// ```
 #[derive(Default, Debug, Clone)]
-pub struct Compiler {
+pub struct Compiler<'a> {
     /// The location of the files to compile
     pub sources: Arc<Vec<String>>,
     /// The output location
     pub output: Option<String>,
     /// Constructor Input Arguments
     pub construct_args: Option<Vec<String>>,
+    /// Constant Overrides
+    pub constant_overrides: Option<BTreeMap<&'a str, Literal>>,
     /// Whether to optimize compilation or not.
     pub optimize: bool,
     /// Generate and log bytecode
@@ -62,19 +66,28 @@ pub struct Compiler {
     pub cached: bool,
 }
 
-impl<'a> Compiler {
+impl<'a> Compiler<'a> {
     /// Public associated function to instantiate a new compiler.
     pub fn new(
         sources: Arc<Vec<String>>,
         output: Option<String>,
         construct_args: Option<Vec<String>>,
+        constant_overrides: Option<BTreeMap<&'a str, Literal>>,
         verbose: bool,
         cached: bool,
     ) -> Self {
         if cfg!(feature = "verbose") || verbose {
             Compiler::init_tracing_subscriber(Some(vec![tracing::Level::INFO.into()]));
         }
-        Self { sources, output, construct_args, optimize: false, bytecode: false, cached }
+        Self {
+            sources,
+            output,
+            construct_args,
+            constant_overrides,
+            optimize: false,
+            bytecode: false,
+            cached,
+        }
     }
 
     /// Tracing
@@ -230,6 +243,7 @@ impl<'a> Compiler {
         let parse_res = parser.parse().map_err(CompilerError::ParserError);
         let mut contract = parse_res?;
         contract.derive_storage_pointers();
+        contract.add_override_constants(&self.constant_overrides);
         tracing::info!(target: "core", "PARSED CONTRACT [{}]", file.path);
 
         // Primary Bytecode Generation
diff --git a/huff_parser/README.md b/huff_parser/README.md
index 1258ced..dc42b52 100644
--- a/huff_parser/README.md
+++ b/huff_parser/README.md
@@ -19,6 +19,8 @@ definition.
 use huff_utils::prelude::*;
 use huff_lexer::{Lexer};
 use huff_parser::{Parser};
+use std::cell::RefCell;
+use std::rc::Rc;
 
 // Create a Lexer from the source code
 let source = "#define macro HELLO_WORLD() = takes(0) returns(0) {}";
@@ -50,7 +52,7 @@ let expected_contract = Contract {
   ],
   invocations: vec![],
   imports: vec![],
-  constants: vec![],
+  constants: Rc::new(RefCell::new(vec![])),
   functions: vec![],
   events: vec![],
   tables: vec![],
diff --git a/huff_parser/src/lib.rs b/huff_parser/src/lib.rs
index 5101a0a..7dcd1a0 100644
--- a/huff_parser/src/lib.rs
+++ b/huff_parser/src/lib.rs
@@ -86,7 +86,7 @@ impl Parser {
                 TokenKind::Constant => {
                     let c = self.parse_constant()?;
                     tracing::info!(target: "parser", "SUCCESSFULLY PARSED CONSTANT {}", c.name);
-                    contract.constants.push(c);
+                    contract.constants.borrow_mut().push(c);
                 }
                 TokenKind::Macro | TokenKind::Fn => {
                     let m = self.parse_macro()?;
diff --git a/huff_utils/src/abi.rs b/huff_utils/src/abi.rs
index e38d76b..72a88c0 100644
--- a/huff_utils/src/abi.rs
+++ b/huff_utils/src/abi.rs
@@ -9,6 +9,8 @@
 //! instance like so:
 //!
 //! ```rust
+//! use std::cell::RefCell;
+//! use std::rc::Rc;
 //! use huff_utils::prelude::*;
 //!
 //! // Generate a default contract for demonstrative purposes.
@@ -17,7 +19,7 @@
 //!     macros: vec![],
 //!     invocations: vec![],
 //!     imports: vec![],
-//!     constants: vec![],
+//!     constants: Rc::new(RefCell::new(vec![])),
 //!     functions: vec![huff_utils::ast::Function {
 //!         name: "CONSTRUCTOR".to_string(),
 //!         signature: [0u8, 0u8, 0u8, 0u8],
diff --git a/huff_utils/src/ast.rs b/huff_utils/src/ast.rs
index 70ee54b..b1d849e 100644
--- a/huff_utils/src/ast.rs
+++ b/huff_utils/src/ast.rs
@@ -9,9 +9,11 @@ use crate::{
     prelude::{Span, TokenKind},
 };
 use std::{
+    cell::RefCell,
     collections::BTreeMap,
     fmt::{Display, Formatter},
     path::PathBuf,
+    rc::Rc,
 };
 
 /// A contained literal
@@ -96,7 +98,7 @@ pub struct Contract {
     /// File Imports
     pub imports: Vec<FilePath>,
     /// Constants
-    pub constants: Vec<ConstantDefinition>,
+    pub constants: Rc<RefCell<Vec<ConstantDefinition>>>,
     /// Functions
     pub functions: Vec<Function>,
     /// Events
@@ -159,7 +161,7 @@ impl Contract {
         tracing::debug!(target: "ast", "ALL AST CONSTANTS: {:?}", storage_pointers);
 
         // Set all the constants to their new values
-        for c in &mut self.constants {
+        for c in self.constants.borrow_mut().iter_mut() {
             match storage_pointers
                 .iter()
                 .filter(|pointer| pointer.0.eq(&c.name))
@@ -217,6 +219,7 @@ impl Contract {
                         // Get the associated constant
                         match self
                             .constants
+                            .borrow()
                             .iter()
                             .filter(|c| c.name.eq(const_name))
                             .collect::<Vec<&ConstantDefinition>>()
@@ -291,6 +294,29 @@ impl Contract {
         //     self.recurse_ast_constants(next_md, storage_pointers, last_p, macros_to_recurse);
         // }
     }
+
+    /// Add override constants to the AST
+    ///
+    /// ## Overview
+    ///
+    /// For each override constant, add it to the AST if it doesn't already exist. Override
+    /// constants can be passed in via the CLI.
+    pub fn add_override_constants(&self, override_constants: &Option<BTreeMap<&str, Literal>>) {
+        if let Some(override_constants) = override_constants {
+            for (name, value) in override_constants {
+                let mut constants = self.constants.borrow_mut();
+                if let Some(c) = constants.iter_mut().find(|c| c.name.as_str().eq(*name)) {
+                    c.value = ConstVal::Literal(*value);
+                } else {
+                    constants.push(ConstantDefinition {
+                        name: name.to_string(),
+                        value: ConstVal::Literal(*value),
+                        span: AstSpan::default(),
+                    });
+                }
+            }
+        }
+    }
 }
 
 /// A function, event, or macro argument


Test patch:
diff --git a/huff_codegen/tests/abigen.rs b/huff_codegen/tests/abigen.rs
index b05bc6a..4c30347 100644
--- a/huff_codegen/tests/abigen.rs
+++ b/huff_codegen/tests/abigen.rs
@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
 
 use huff_codegen::Codegen;
 use huff_utils::{ast, prelude::*};
@@ -18,7 +18,7 @@ fn constructs_valid_abi() {
         macros: vec![constructor],
         invocations: vec![],
         imports: vec![],
-        constants: vec![],
+        constants: Rc::new(RefCell::new(vec![])),
         functions: vec![],
         events: vec![],
         tables: vec![],
@@ -41,7 +41,7 @@ fn constructs_valid_abi() {
 
 #[test]
 fn missing_constructor_fails() {
-    let _constructor = ast::MacroDefinition {
+    let _constructor = MacroDefinition {
         name: "CONSTRUCTOR".to_string(),
         parameters: vec![],
         statements: vec![],
@@ -54,7 +54,7 @@ fn missing_constructor_fails() {
         macros: vec![],
         invocations: vec![],
         imports: vec![],
-        constants: vec![],
+        constants: Rc::new(RefCell::new(vec![])),
         functions: vec![],
         events: vec![],
         tables: vec![],
diff --git a/huff_core/tests/file_resolution.rs b/huff_core/tests/file_resolution.rs
index 507ed36..7b215e0 100644
--- a/huff_core/tests/file_resolution.rs
+++ b/huff_core/tests/file_resolution.rs
@@ -5,7 +5,7 @@ use huff_utils::prelude::{CompilerError, OutputLocation, UnpackError};
 
 #[test]
 fn test_get_outputs_no_output() {
-    let compiler: Compiler = Compiler::new(Arc::new(vec![]), None, None, false, false);
+    let compiler: Compiler = Compiler::new(Arc::new(vec![]), None, None, None, false, false);
     let ol: OutputLocation = compiler.get_outputs();
     assert_eq!(ol, OutputLocation::default());
 }
@@ -13,7 +13,7 @@ fn test_get_outputs_no_output() {
 #[test]
 fn test_get_outputs_with_output() {
     let compiler: Compiler =
-        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, false, false);
+        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, None, false, false);
     let ol: OutputLocation = compiler.get_outputs();
     assert_eq!(ol, OutputLocation("./test_out/".to_string()));
 }
@@ -21,7 +21,7 @@ fn test_get_outputs_with_output() {
 #[test]
 fn test_transform_paths() {
     let _compiler: Compiler =
-        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, false, false);
+        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, None, false, false);
     let path_bufs: Result<Vec<PathBuf>, CompilerError<'_>> = Compiler::transform_paths(&vec![
         "../huff-examples/erc20/contracts/ERC20.huff".to_string(),
         "../huff-examples/erc20/contracts/utils/".to_string(),
@@ -55,7 +55,7 @@ fn test_transform_paths() {
 #[test]
 fn test_transform_paths_non_huff() {
     let _compiler: Compiler =
-        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, false, false);
+        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, None, false, false);
     let path_bufs: Result<Vec<PathBuf>, CompilerError<'_>> =
         Compiler::transform_paths(&vec!["./ERC20.txt".to_string()]);
     assert!(path_bufs.is_err());
@@ -72,7 +72,7 @@ fn test_transform_paths_non_huff() {
 #[test]
 fn test_transform_paths_no_dir() {
     let _compiler: Compiler =
-        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, false, false);
+        Compiler::new(Arc::new(vec![]), Some("./test_out/".to_string()), None, None, false, false);
     let path_bufs: Result<Vec<PathBuf>, CompilerError<'_>> =
         Compiler::transform_paths(&vec!["./examples/random_dir/".to_string()]);
     assert!(path_bufs.is_err());
diff --git a/huff_core/tests/gen_artifact.rs b/huff_core/tests/gen_artifact.rs
index a1d4e82..eab2f3b 100644
--- a/huff_core/tests/gen_artifact.rs
+++ b/huff_core/tests/gen_artifact.rs
@@ -31,7 +31,7 @@ fn test_missing_constructor() {
     };
 
     // Instantiate a new compiler
-    let compiler = Compiler::new(Arc::new(vec![]), None, None, false, false);
+    let compiler = Compiler::new(Arc::new(vec![]), None, None, None, false, false);
 
     // Generate the compile artifact
     let arc_source = Arc::new(full_source);
@@ -76,7 +76,8 @@ fn test_missing_constructor_with_inputs() {
     };
 
     // Instantiate a new compiler
-    let compiler = Compiler::new(Arc::new(vec![]), None, Some(vec!["0".to_string()]), false, false);
+    let compiler =
+        Compiler::new(Arc::new(vec![]), None, Some(vec!["0".to_string()]), None, false, false);
 
     // Generate the compile artifact
     let arc_source = Arc::new(full_source);
diff --git a/huff_parser/tests/constant.rs b/huff_parser/tests/constant.rs
index 4653438..e7f4126 100644
--- a/huff_parser/tests/constant.rs
+++ b/huff_parser/tests/constant.rs
@@ -12,7 +12,7 @@ fn test_parses_free_storage_pointer_constant() {
     let contract = parser.parse().unwrap();
     assert_eq!(parser.current_token.kind, TokenKind::Eof);
 
-    let fsp_constant = contract.constants[0].clone();
+    let fsp_constant = contract.constants.borrow()[0].clone();
     assert_eq!(
         fsp_constant,
         ConstantDefinition {
@@ -44,7 +44,7 @@ fn test_parses_literal_constant() {
         str_to_bytes32("8C5BE1E5EBEC7D5BD14F71427D1E84F3DD0314C0F7B2291E5B200AC8C7C3B925");
 
     // Check Literal
-    let fsp_constant = contract.constants[0].clone();
+    let fsp_constant = contract.constants.borrow()[0].clone();
     assert_eq!(
         fsp_constant,
         ConstantDefinition {
diff --git a/huff_parser/tests/storage_pointer_derivation.rs b/huff_parser/tests/storage_pointer_derivation.rs
index c3598e4..f1769a3 100644
--- a/huff_parser/tests/storage_pointer_derivation.rs
+++ b/huff_parser/tests/storage_pointer_derivation.rs
@@ -15,7 +15,7 @@ fn derives_storage_pointers() {
     assert_eq!(parser.current_token.kind, TokenKind::Eof);
 
     // Ensure that the constant definitions were parsed correctly
-    let fsp_constant = contract.constants[0].clone();
+    let fsp_constant = contract.constants.borrow()[0].clone();
     assert_eq!(
         fsp_constant,
         ConstantDefinition {
@@ -31,7 +31,7 @@ fn derives_storage_pointers() {
         }
     );
 
-    let fsp_constant = contract.constants[1].clone();
+    let fsp_constant = contract.constants.borrow()[1].clone();
     assert_eq!(
         fsp_constant,
         ConstantDefinition {
@@ -47,7 +47,7 @@ fn derives_storage_pointers() {
         }
     );
 
-    let num_constant = contract.constants[2].clone();
+    let num_constant = contract.constants.borrow()[2].clone();
     assert_eq!(
         num_constant,
         ConstantDefinition {
@@ -67,7 +67,13 @@ fn derives_storage_pointers() {
     contract.derive_storage_pointers();
 
     // Ensure that the storage pointers were set for the FSP constants in the AST
-    assert_eq!(contract.constants[0].value, ConstVal::FreeStoragePointer(FreeStoragePointer));
-    assert_eq!(contract.constants[1].value, ConstVal::FreeStoragePointer(FreeStoragePointer));
-    assert_eq!(contract.constants[2].value, ConstVal::Literal(str_to_bytes32("a57B")));
+    assert_eq!(
+        contract.constants.borrow()[0].value,
+        ConstVal::FreeStoragePointer(FreeStoragePointer)
+    );
+    assert_eq!(
+        contract.constants.borrow()[1].value,
+        ConstVal::FreeStoragePointer(FreeStoragePointer)
+    );
+    assert_eq!(contract.constants.borrow()[2].value, ConstVal::Literal(str_to_bytes32("a57B")));
 }