Repository: huff-language/huff-rs
Pull request #174 (2022-07-26 15:22:34)

feat(huff_codegen): Custom Errors
# Overview

Adds custom error definitions and a builtin to push their left-padded selectors to the stack.

Syntax:
```
#define error MyCustomError(uint256)

#define macro MAIN() = takes (0) returns (0) {
    __ERROR(MyCustomError)  // [left_padded_error_selector]
    0x00 mstore             // []
    0x01 0x04 mstore        // []
    0x24 0x00 revert
}
```

## Checklist
- [x] Lex/parse error definitions
  - [x] Tests 
- [x] Add left-padded error selector builtin
  - [x] Tests 
- [x] Pull out hashing logic into `bytes_util`, it's used in more than a couple places now
- [x] Clean up builtin detection in `huff_lexer`
  - ![image](https://user-images.githubusercontent.com/8406232/181054696-6eedcb4e-799c-46e6-ad51-b89789a7a4ed.png) 


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index e749927..462070c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -671,7 +671,6 @@ dependencies = [
  "huff_utils",
  "proptest",
  "serde_json",
- "tiny-keccak",
  "tracing",
  "uuid",
 ]
@@ -716,7 +715,6 @@ dependencies = [
  "huff_utils",
  "proptest",
  "regex",
- "tiny-keccak",
  "tracing",
 ]
 
@@ -735,6 +733,7 @@ dependencies = [
  "serde_json",
  "strum",
  "strum_macros",
+ "tiny-keccak",
  "tracing",
  "uuid",
 ]
diff --git a/huff_codegen/Cargo.toml b/huff_codegen/Cargo.toml
index 67dab54..0bc30f3 100644
--- a/huff_codegen/Cargo.toml
+++ b/huff_codegen/Cargo.toml
@@ -19,4 +19,3 @@ ethers-core = "0.13.0"
 hex = "0.4.3"
 tracing = "0.1.34"
 uuid = { version = "1.1.1", features = ["v4"] }
-tiny-keccak = { version = "2.0.2", features = ["keccak"] }
diff --git a/huff_codegen/README.md b/huff_codegen/README.md
index 4a63122..7f33a4b 100644
--- a/huff_codegen/README.md
+++ b/huff_codegen/README.md
@@ -91,6 +91,7 @@ let contract = Contract {
   invocations: vec![],
   imports: vec![],
   constants: Rc::new(RefCell::new(vec![])),
+  errors: vec![],
   functions: vec![],
   events: vec![],
   tables: vec![],
@@ -145,6 +146,7 @@ let contract = Contract {
   invocations: vec![],
   imports: vec![],
   constants: Rc::new(RefCell::new(vec![])),
+  errors: vec![],
   functions: vec![],
   events: vec![],
   tables: vec![],
diff --git a/huff_codegen/src/irgen/statements.rs b/huff_codegen/src/irgen/statements.rs
index 6e85671..037105e 100644
--- a/huff_codegen/src/irgen/statements.rs
+++ b/huff_codegen/src/irgen/statements.rs
@@ -1,5 +1,4 @@
 use huff_utils::prelude::*;
-use tiny_keccak::{Hasher, Keccak};
 
 use crate::Codegen;
 
@@ -287,9 +286,7 @@ pub fn statement_gen(
                         bytes.push((starting_offset, Bytes(push_bytes)));
                     } else if let Some(s) = &bf.args[0].name {
                         let mut signature = [0u8; 4]; // Only keep first 4 bytes
-                        let mut hasher = Keccak::v256();
-                        hasher.update(s.as_bytes());
-                        hasher.finalize(&mut signature);
+                        hash_bytes(&mut signature, s);
 
                         let sig = hex::encode(signature);
                         let push_bytes = format!("{:02x}{}", 95 + sig.len() / 2, sig);
@@ -340,9 +337,7 @@ pub fn statement_gen(
                         bytes.push((starting_offset, Bytes(push_bytes)));
                     } else if let Some(s) = &bf.args[0].name {
                         let mut hash = [0u8; 32];
-                        let mut hasher = Keccak::v256();
-                        hasher.update(s.as_bytes());
-                        hasher.finalize(&mut hash);
+                        hash_bytes(&mut hash, s);
 
                         let hash = hex::encode(hash);
                         let push_bytes = format!("{:02x}{}", 95 + hash.len() / 2, hash);
@@ -363,6 +358,49 @@ pub fn statement_gen(
                         })
                     }
                 }
+                BuiltinFunctionKind::Error => {
+                    if bf.args.len() != 1 {
+                        tracing::error!(
+                            target: "codegen",
+                            "Incorrect number of arguments passed to __ERROR, should be 1: {}",
+                            bf.args.len()
+                        );
+                        return Err(CodegenError {
+                            kind: CodegenErrorKind::InvalidArguments(format!(
+                                "Incorrect number of arguments passed to __ERROR, should be 1: {}",
+                                bf.args.len()
+                            )),
+                            span: bf.span.clone(),
+                            token: None,
+                        })
+                    }
+
+                    if let Some(error) = contract
+                        .errors
+                        .iter()
+                        .find(|e| bf.args[0].name.as_ref().unwrap().eq(&e.name))
+                    {
+                        // Add 28 bytes to left-pad the 4 byte selector
+                        let selector =
+                            format!("{}{}", hex::encode(error.selector), "00".repeat(28));
+                        let push_bytes = format!("{:02x}{}", 95 + selector.len() / 2, selector);
+                        *offset += push_bytes.len() / 2;
+                        bytes.push((starting_offset, Bytes(push_bytes)));
+                    } else {
+                        tracing::error!(
+                            target: "codegen",
+                            "MISSING ERROR DEFINITION PASSED TO __ERROR: \"{}\"",
+                            bf.args[0].name.as_ref().unwrap()
+                        );
+                        return Err(CodegenError {
+                            kind: CodegenErrorKind::MissingErrorDefinition(
+                                bf.args[0].name.as_ref().unwrap().to_string(),
+                            ),
+                            span: bf.span.clone(),
+                            token: None,
+                        })
+                    }
+                }
             }
         }
         sty => {
diff --git a/huff_lexer/src/lib.rs b/huff_lexer/src/lib.rs
index 1640a27..7d93175 100644
--- a/huff_lexer/src/lib.rs
+++ b/huff_lexer/src/lib.rs
@@ -281,6 +281,7 @@ impl<'a> Lexer<'a> {
             Some(TokenKind::Fn) |
             Some(TokenKind::Function) |
             Some(TokenKind::Constant) |
+            Some(TokenKind::Error) |
             Some(TokenKind::Event) |
             Some(TokenKind::JumpTable) |
             Some(TokenKind::JumpTablePacked) |
@@ -383,6 +384,7 @@ impl<'a> Iterator for Lexer<'a> {
                         TokenKind::Fn,
                         TokenKind::Function,
                         TokenKind::Constant,
+                        TokenKind::Error,
                         TokenKind::Takes,
                         TokenKind::Returns,
                         TokenKind::Event,
@@ -424,7 +426,9 @@ impl<'a> Iterator for Lexer<'a> {
                             TokenKind::Macro | TokenKind::Fn => {
                                 self.context = Context::MacroDefinition
                             }
-                            TokenKind::Function | TokenKind::Event => self.context = Context::Abi,
+                            TokenKind::Function | TokenKind::Event | TokenKind::Error => {
+                                self.context = Context::Abi
+                            }
                             TokenKind::Constant => self.context = Context::Constant,
                             TokenKind::CodeTable => self.context = Context::CodeTableBody,
                             _ => (),
@@ -557,14 +561,7 @@ impl<'a> Iterator for Lexer<'a> {
                         let slice = self.slice();
                         // Check for built-in function calls
                         if self.context == Context::MacroBody &&
-                            matches!(
-                                slice.as_ref(),
-                                "__codesize" |
-                                    "__tablesize" |
-                                    "__tablestart" |
-                                    "__FUNC_SIG" |
-                                    "__EVENT_HASH" /* TODO: Clean this process up */
-                            )
+                            BuiltinFunctionKind::try_from(&slice).is_ok()
                         {
                             TokenKind::BuiltinFunction(slice)
                         } else {
diff --git a/huff_parser/Cargo.toml b/huff_parser/Cargo.toml
index 8706b56..0b41bc0 100644
--- a/huff_parser/Cargo.toml
+++ b/huff_parser/Cargo.toml
@@ -15,7 +15,6 @@ keywords = ["huff", "rust", "evm", "bytecode", "compiler"]
 huff_utils = { path = "../huff_utils" }
 huff_lexer = { path = "../huff_lexer" }
 proptest = "1.0.0"
-tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 tracing = "0.1.34"
 hex = "0.4.3"
 regex = "1.6.0"
\ No newline at end of file
diff --git a/huff_parser/README.md b/huff_parser/README.md
index dc42b52..dca0fae 100644
--- a/huff_parser/README.md
+++ b/huff_parser/README.md
@@ -53,6 +53,7 @@ let expected_contract = Contract {
   invocations: vec![],
   imports: vec![],
   constants: Rc::new(RefCell::new(vec![])),
+  errors: vec![],
   functions: vec![],
   events: vec![],
   tables: vec![],
diff --git a/huff_parser/src/lib.rs b/huff_parser/src/lib.rs
index 1f5eb0c..a8c65d9 100644
--- a/huff_parser/src/lib.rs
+++ b/huff_parser/src/lib.rs
@@ -7,13 +7,12 @@
 use huff_utils::{
     ast::*,
     error::*,
-    prelude::{str_to_bytes32, FileSource, Span},
+    prelude::{hash_bytes, str_to_bytes32, FileSource, Span},
     token::{Token, TokenKind},
     types::*,
 };
 use regex::Regex;
 use std::path::Path;
-use tiny_keccak::{Hasher, Keccak};
 
 /// The Parser
 #[derive(Debug, Clone)]
@@ -88,6 +87,11 @@ impl Parser {
                         tracing::info!(target: "parser", "SUCCESSFULLY PARSED CONSTANT {}", c.name);
                         contract.constants.borrow_mut().push(c);
                     }
+                    TokenKind::Error => {
+                        let e = self.parse_custom_error()?;
+                        tracing::info!(target: "parser", "SUCCESSFULLY PARSED ERROR {}", e.name);
+                        contract.errors.push(e);
+                    }
                     TokenKind::Macro | TokenKind::Fn => {
                         let m = self.parse_macro()?;
                         tracing::info!(target: "parser", "SUCCESSFULLY PARSED MACRO {}", m.name);
@@ -255,7 +259,7 @@ impl Parser {
         };
 
         // function inputs should be next
-        let inputs: Vec<Argument> = self.parse_args(true, true, false, false)?;
+        let inputs = self.parse_args(true, true, false, false)?;
         // function type should be next
         let fn_type = match self.current_token.kind.clone() {
             TokenKind::View => FunctionType::View,
@@ -278,14 +282,12 @@ impl Parser {
         // next token should be of `TokenKind::Returns`
         self.match_kind(TokenKind::Returns)?;
         // function outputs should be next
-        let outputs: Vec<Argument> = self.parse_args(true, true, false, false)?;
+        let outputs = self.parse_args(true, true, false, false)?;
 
         let mut signature = [0u8; 4]; // Only keep first 4 bytes
-        let mut hasher = Keccak::v256();
         let input_types =
             inputs.iter().map(|i| i.arg_type.as_ref().unwrap().clone()).collect::<Vec<_>>();
-        hasher.update(format!("{}({})", name, input_types.join(",")).as_bytes());
-        hasher.finalize(&mut signature);
+        hash_bytes(&mut signature, &format!("{}({})", name, input_types.join(",")));
 
         Ok(Function {
             name,
@@ -319,14 +321,12 @@ impl Parser {
         };
 
         // Parse the event's parameters
-        let parameters: Vec<Argument> = self.parse_args(true, true, true, false)?;
+        let parameters = self.parse_args(true, true, true, false)?;
 
         let mut hash = [0u8; 32];
-        let mut hasher = Keccak::v256();
         let input_types =
             parameters.iter().map(|i| i.arg_type.as_ref().unwrap().clone()).collect::<Vec<_>>();
-        hasher.update(format!("{}({})", name, input_types.join(",")).as_bytes());
-        hasher.finalize(&mut hash);
+        hash_bytes(&mut hash, &format!("{}({})", name, input_types.join(",")));
 
         Ok(Event { name, parameters, span: AstSpan(self.spans.clone()), hash })
     }
@@ -343,12 +343,10 @@ impl Parser {
             TokenKind::Ident(const_name) => const_name,
             _ => {
                 tracing::error!(target: "parser", "TOKEN MISMATCH - EXPECTED IDENT, GOT: {}", tok);
-                let new_spans = self.spans.clone();
-                self.spans = vec![];
                 return Err(ParserError {
                     kind: ParserErrorKind::UnexpectedType(tok),
                     hint: Some("Expected constant name.".to_string()),
-                    spans: AstSpan(new_spans),
+                    spans: AstSpan(self.spans.clone()),
                 })
             }
         };
@@ -386,6 +384,41 @@ impl Parser {
         Ok(ConstantDefinition { name, value, span: AstSpan(new_spans) })
     }
 
+    /// Parse a custom error definition.
+    pub fn parse_custom_error(&mut self) -> Result<ErrorDefinition, ParserError> {
+        // Error Identifier
+        self.match_kind(TokenKind::Error)?;
+
+        // Parse the error name
+        self.match_kind(TokenKind::Ident("x".to_string()))?;
+        let tok = self.peek_behind().unwrap().kind;
+        let name = match tok {
+            TokenKind::Ident(err_name) => err_name,
+            _ => {
+                tracing::error!(target: "parser", "TOKEN MISMATCH - EXPECTED IDENT, GOT: {}", tok);
+                return Err(ParserError {
+                    kind: ParserErrorKind::UnexpectedType(tok),
+                    hint: Some("Expected error name.".to_string()),
+                    spans: AstSpan(self.spans.clone()),
+                })
+            }
+        };
+
+        // Get arguments for signature
+        let inputs = self.parse_args(true, true, false, false)?;
+
+        let mut selector = [0u8; 4]; // Only keep first 4 bytes
+        let input_types =
+            inputs.iter().map(|i| i.arg_type.as_ref().unwrap().clone()).collect::<Vec<_>>();
+        hash_bytes(&mut selector, &format!("{}({})", name, input_types.join(",")));
+
+        // Clone spans and set to nothing
+        let new_spans = self.spans.clone();
+        self.spans = vec![];
+
+        Ok(ErrorDefinition { name, selector, span: AstSpan(new_spans) })
+    }
+
     /// Parses a macro.
     ///
     /// It should parse the following : macro MACRO_NAME(args...) = takes (x) returns (n) {...}
@@ -396,7 +429,7 @@ impl Parser {
             self.match_kind(TokenKind::Ident("MACRO_NAME".to_string()))?.to_string();
         tracing::info!(target: "parser", "PARSING MACRO: \"{}\"", macro_name);
 
-        let macro_arguments: Vec<Argument> = self.parse_args(true, false, false, false)?;
+        let macro_arguments = self.parse_args(true, false, false, false)?;
         self.match_kind(TokenKind::Assign)?;
         self.match_kind(TokenKind::Takes)?;
         let macro_takes: usize = self.parse_single_arg()?;
@@ -512,7 +545,7 @@ impl Parser {
                     tracing::info!(target: "parser", "PARSING MACRO BODY: [BUILTIN FN: {}({:?})]", f, args);
                     statements.push(Statement {
                         ty: StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
-                            kind: BuiltinFunctionKind::from(f.as_str()),
+                            kind: BuiltinFunctionKind::from(f),
                             args,
                             span: AstSpan(curr_spans.clone()),
                         }),
diff --git a/huff_utils/Cargo.toml b/huff_utils/Cargo.toml
index 202cfd4..5f69218 100644
--- a/huff_utils/Cargo.toml
+++ b/huff_utils/Cargo.toml
@@ -25,3 +25,4 @@ uuid = { version = "1.1.1", features = ["v4"] }
 pathdiff = "0.2.1"
 ethers-core = "0.13.0"
 itertools = "0.10.3"
+tiny-keccak = { version = "2.0.2", features = ["keccak"] }
\ No newline at end of file
diff --git a/huff_utils/src/abi.rs b/huff_utils/src/abi.rs
index 72a88c0..bc820d9 100644
--- a/huff_utils/src/abi.rs
+++ b/huff_utils/src/abi.rs
@@ -20,6 +20,7 @@
 //!     invocations: vec![],
 //!     imports: vec![],
 //!     constants: Rc::new(RefCell::new(vec![])),
+//!     errors: 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 b1d849e..48f23ed 100644
--- a/huff_utils/src/ast.rs
+++ b/huff_utils/src/ast.rs
@@ -12,6 +12,7 @@ use std::{
     cell::RefCell,
     collections::BTreeMap,
     fmt::{Display, Formatter},
+    panic,
     path::PathBuf,
     rc::Rc,
 };
@@ -99,6 +100,8 @@ pub struct Contract {
     pub imports: Vec<FilePath>,
     /// Constants
     pub constants: Rc<RefCell<Vec<ConstantDefinition>>>,
+    /// Custom Errors
+    pub errors: Vec<ErrorDefinition>,
     /// Functions
     pub functions: Vec<Function>,
     /// Events
@@ -621,6 +624,17 @@ pub struct ConstantDefinition {
     pub span: AstSpan,
 }
 
+/// An Error Definition
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub struct ErrorDefinition {
+    /// The Error name
+    pub name: String,
+    /// The Error's selector
+    pub selector: [u8; 4],
+    /// The Span of the Constant Definition
+    pub span: AstSpan,
+}
+
 /// A Jump Destination
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub struct Label {
@@ -639,7 +653,7 @@ pub struct BuiltinFunctionCall {
     pub kind: BuiltinFunctionKind,
     /// Arguments for the builtin function call.
     /// TODO: Maybe make a better type for this other than `Argument`? Would be nice if it pointed
-    /// directly to the macro/table.
+    ///       directly to the macro/table.
     pub args: Vec<Argument>,
     /// The builtin function call span
     pub span: AstSpan,
@@ -658,17 +672,34 @@ pub enum BuiltinFunctionKind {
     FunctionSignature,
     /// Event hash function
     EventHash,
+    /// Error selector function
+    Error,
 }
 
-impl From<&str> for BuiltinFunctionKind {
-    fn from(s: &str) -> Self {
-        match s {
+impl From<String> for BuiltinFunctionKind {
+    fn from(value: String) -> Self {
+        match value.as_str() {
             "__tablesize" => BuiltinFunctionKind::Tablesize,
             "__codesize" => BuiltinFunctionKind::Codesize,
             "__tablestart" => BuiltinFunctionKind::Tablestart,
             "__FUNC_SIG" => BuiltinFunctionKind::FunctionSignature,
             "__EVENT_HASH" => BuiltinFunctionKind::EventHash,
-            _ => panic!("Invalid Builtin Function Kind"), // TODO: Better error handling
+            "__ERROR" => BuiltinFunctionKind::Error,
+            _ => panic!("Invalid Builtin Function Kind"), /* This should never be reached,
+                                                           * builtins are validated with a
+                                                           * `try_from` call in the lexer. */
+        }
+    }
+}
+
+impl TryFrom<&String> for BuiltinFunctionKind {
+    type Error = ();
+
+    fn try_from(value: &String) -> Result<Self, <BuiltinFunctionKind as TryFrom<&String>>::Error> {
+        if let Ok(kind) = panic::catch_unwind(|| BuiltinFunctionKind::from(value.to_owned())) {
+            Ok(kind)
+        } else {
+            Err(())
         }
     }
 }
diff --git a/huff_utils/src/bytes_util.rs b/huff_utils/src/bytes_util.rs
index d7b21ac..2018ecd 100644
--- a/huff_utils/src/bytes_util.rs
+++ b/huff_utils/src/bytes_util.rs
@@ -1,3 +1,5 @@
+use tiny_keccak::{Hasher, Keccak};
+
 /// Convert a string slice to a `[u8; 32]`
 /// Pads zeros to the left of significant bytes in the `[u8; 32]` slice.
 /// i.e. 0xa57b becomes `[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -52,3 +54,10 @@ pub fn str_to_vec(s: &str) -> Result<Vec<u8>, std::num::ParseIntError> {
         (0..s.len()).step_by(2).map(|c| u8::from_str_radix(&s[c..c + 2], 16)).collect();
     bytes
 }
+
+/// Hash a string with Keccak256
+pub fn hash_bytes(dest: &mut [u8], to_hash: &String) {
+    let mut hasher = Keccak::v256();
+    hasher.update(to_hash.as_bytes());
+    hasher.finalize(dest);
+}
diff --git a/huff_utils/src/error.rs b/huff_utils/src/error.rs
index f24de3c..67d9309 100644
--- a/huff_utils/src/error.rs
+++ b/huff_utils/src/error.rs
@@ -141,6 +141,8 @@ pub enum CodegenErrorKind {
     MissingEventInterface(String),
     /// Missing Constant Definition
     MissingConstantDefinition(String),
+    /// Missing Error Definition
+    MissingErrorDefinition(String),
     /// Abi Generation Failure
     AbiGenerationFailure,
     /// Unmatched Jump
@@ -191,6 +193,9 @@ impl<W: Write> Report<W> for CodegenError {
             CodegenErrorKind::MissingConstantDefinition(cd) => {
                 write!(f.out, "Missing Constant Definition for \"{}\"!", cd)
             }
+            CodegenErrorKind::MissingErrorDefinition(ed) => {
+                write!(f.out, "Missing Error Definition for \"{}\"!", ed)
+            }
             CodegenErrorKind::AbiGenerationFailure => write!(f.out, "Abi generation failure!"),
             CodegenErrorKind::UnmatchedJumpLabel => write!(f.out, "Unmatched jump label!"),
             CodegenErrorKind::IOError(ioe) => write!(f.out, "IO ERROR: {:?}", ioe),
@@ -473,6 +478,9 @@ impl<'a> fmt::Display for CompilerError<'a> {
                 CodegenErrorKind::MissingConstantDefinition(_) => {
                     write!(f, "\nError: Missing Constant Definition\n{}\n", ce.span.error(None))
                 }
+                CodegenErrorKind::MissingErrorDefinition(_) => {
+                    write!(f, "\nError: Missing Error Definition\n{}\n", ce.span.error(None))
+                }
                 CodegenErrorKind::AbiGenerationFailure => {
                     write!(f, "\nError: ABI Generation Failed\n{}\n", ce.span.error(None))
                 }
diff --git a/huff_utils/src/token.rs b/huff_utils/src/token.rs
index c572a7b..a38e616 100644
--- a/huff_utils/src/token.rs
+++ b/huff_utils/src/token.rs
@@ -43,6 +43,8 @@ pub enum TokenKind {
     Event,
     /// "constant" keyword
     Constant,
+    /// "error" keyword
+    Error,
     /// "takes" keyword
     Takes,
     /// "returns" keyword
@@ -132,6 +134,7 @@ impl fmt::Display for TokenKind {
             TokenKind::Function => "function",
             TokenKind::Event => "event",
             TokenKind::Constant => "constant",
+            TokenKind::Error => "error",
             TokenKind::View => "view",
             TokenKind::Pure => "pure",
             TokenKind::Payable => "payable",


Test patch:
diff --git a/huff_codegen/tests/abigen.rs b/huff_codegen/tests/abigen.rs
index 4c30347..1761e30 100644
--- a/huff_codegen/tests/abigen.rs
+++ b/huff_codegen/tests/abigen.rs
@@ -19,6 +19,7 @@ fn constructs_valid_abi() {
         invocations: vec![],
         imports: vec![],
         constants: Rc::new(RefCell::new(vec![])),
+        errors: vec![],
         functions: vec![],
         events: vec![],
         tables: vec![],
@@ -55,6 +56,7 @@ fn missing_constructor_fails() {
         invocations: vec![],
         imports: vec![],
         constants: Rc::new(RefCell::new(vec![])),
+        errors: vec![],
         functions: vec![],
         events: vec![],
         tables: vec![],
diff --git a/huff_core/tests/builtins.rs b/huff_core/tests/builtins.rs
index 9c5914c..55592c8 100644
--- a/huff_core/tests/builtins.rs
+++ b/huff_core/tests/builtins.rs
@@ -477,3 +477,74 @@ fn test_event_hash_builtin() {
         String::from("7fbeabacc8ffedac16e9a60acdb2ca743d80c2ebb44977a93fa8e483c74d2b35a87fbeabacc8ffedac16e9a60acdb2ca743d80c2ebb44977a93fa8e483c74d2b35a87fbeabacc8ffedac16e9a60acdb2ca743d80c2ebb44977a93fa8e483c74d2b35a8600055")
     );
 }
+
+#[test]
+fn test_error_selector_builtin() {
+    let source: &str = r#"
+        // Define our custom error
+        #define error PanicError(uint256 panicCode)
+        #define error Error(string)
+
+        #define macro PANIC() = takes (1) returns (0) {
+            // Input stack:          [panic_code]
+            __ERROR(PanicError)   // [panic_error_selector, panic_code]
+            0x00 mstore           // [panic_code]
+            0x04 mstore           // []
+            0x24 0x00 revert
+        }
+
+        #define macro REQUIRE() = takes (3) returns (0) {
+            // Input stack:          [condition, message_length, message]
+            continue jumpi        // [message]
+
+            __ERROR(Error)        // [error_selector, message_length, message]
+            0x00 mstore           // [message_length, message]
+            0x20 0x04 mstore      // [message_length, message]
+            0x24 mstore           // [message]
+            0x44 mstore           // []
+
+            0x64 0x00 revert
+
+            continue:
+                pop               // []
+        }
+
+        #define macro MAIN() = takes (0) returns (0) {
+            // dummy macro invocations so they're included in the runtime bytecode
+            PANIC()
+            REQUIRE()
+        }
+    "#;
+
+    // Parse tokens
+    let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(flattened_source);
+    let tokens = lexer.into_iter().map(|x| x.unwrap()).collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, None);
+
+    // Parse the AST
+    let mut contract = parser.parse().unwrap();
+
+    // Derive storage pointers
+    contract.derive_storage_pointers();
+
+    // Instantiate Codegen
+    let cg = Codegen::new();
+
+    // The codegen instance should have no artifact
+    assert!(cg.artifact.is_none());
+
+    // Have Codegen create the runtime bytecode
+    let r_bytes = Codegen::generate_main_bytecode(&contract).unwrap();
+    assert_eq!(&r_bytes[2..66], "be20788c00000000000000000000000000000000000000000000000000000000");
+    assert_eq!(
+        &r_bytes[98..162],
+        "08c379a000000000000000000000000000000000000000000000000000000000"
+    );
+    assert_eq!(
+        r_bytes,
+        String::from(
+            "7fbe20788c0000000000000000000000000000000000000000000000000000000060005260045260246000fd610064577f08c379a000000000000000000000000000000000000000000000000000000000600052602060045260245260445260646000fd5b50"
+        )
+    );
+}
diff --git a/huff_lexer/tests/builtins.rs b/huff_lexer/tests/builtins.rs
index cce87e8..d69295c 100644
--- a/huff_lexer/tests/builtins.rs
+++ b/huff_lexer/tests/builtins.rs
@@ -4,7 +4,8 @@ use std::ops::Deref;
 
 #[test]
 fn parses_builtin_function_in_macro_body() {
-    let builtin_funcs = ["__codesize", "__tablesize", "__tablestart", "__FUNC_SIG", "__EVENT_HASH"];
+    let builtin_funcs =
+        ["__codesize", "__tablesize", "__tablestart", "__FUNC_SIG", "__EVENT_HASH", "__ERROR"];
 
     for builtin in builtin_funcs {
         let source = &format!(
@@ -69,7 +70,8 @@ fn parses_builtin_function_in_macro_body() {
 #[test]
 #[should_panic]
 fn fails_to_parse_builtin_outside_macro_body() {
-    let builtin_funcs = ["__codesize", "__tablesize", "__tablestart", "__FUNC_SIG", "__EVENT_HASH"];
+    let builtin_funcs =
+        ["__codesize", "__tablesize", "__tablestart", "__FUNC_SIG", "__EVENT_HASH", "__ERROR"];
 
     for builtin in builtin_funcs {
         let source = &format!("{}(MAIN)", builtin);
@@ -95,3 +97,68 @@ fn fails_to_parse_builtin_outside_macro_body() {
         assert!(lexer.eof);
     }
 }
+
+#[test]
+#[should_panic]
+fn fails_to_parse_invalid_builtin() {
+    let builtin_funcs = ["__not_a_builtin", "__another_not_a_builtin", "__last_not_a_builtin"];
+
+    for builtin in builtin_funcs {
+        let source = &format!(
+            r#"
+            #define macro TEST() = takes(0) returns(0) {}
+                {}(MAIN)
+            {}
+            "#,
+            "{", builtin, "}",
+        );
+        let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+        let mut lexer = Lexer::new(flattened_source.clone());
+        assert_eq!(lexer.source, flattened_source);
+
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // #define
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // macro
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // TEST
+        let _ = lexer.next(); // open parenthesis
+        let _ = lexer.next(); // close parenthesis
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // =
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // takes
+        let _ = lexer.next(); // open parenthesis
+        let _ = lexer.next(); // 0
+        let _ = lexer.next(); // close parenthesis
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // returns
+        let _ = lexer.next(); // open parenthesis
+        let _ = lexer.next(); // 0
+        let _ = lexer.next(); // close parenthesis
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // {
+        let _ = lexer.next(); // whitespace
+
+        // The builtin fn should be parsed as a `TokenKind::BuiltinFunction` here.
+        let tok = lexer.next();
+        let unwrapped = tok.unwrap().unwrap();
+        let builtin_span = Span::new(74..74 + builtin.len(), None);
+        assert_eq!(
+            unwrapped,
+            Token::new(TokenKind::BuiltinFunction(builtin.to_string()), builtin_span.clone())
+        );
+        assert_eq!(lexer.current_span().deref(), &builtin_span);
+
+        let _ = lexer.next(); // open parenthesis
+        let _ = lexer.next(); // MAIN
+        let _ = lexer.next(); // close parenthesis
+        let _ = lexer.next(); // whitespace
+        let _ = lexer.next(); // }
+        let _ = lexer.next(); // whitespace
+
+        // We covered the whole source
+        assert_eq!(lexer.current_span().end, source.len());
+        assert!(lexer.eof);
+    }
+}
diff --git a/huff_lexer/tests/keywords.rs b/huff_lexer/tests/keywords.rs
index 5f9554d..8cabd66 100644
--- a/huff_lexer/tests/keywords.rs
+++ b/huff_lexer/tests/keywords.rs
@@ -305,6 +305,7 @@ fn parses_function_definition_with_keyword_name() {
         "fn",
         "function",
         "constant",
+        "error",
         "takes",
         "returns",
         "define",
@@ -368,6 +369,7 @@ fn parses_label_with_keyword_name() {
         "fn",
         "function",
         "constant",
+        "error",
         "takes",
         "returns",
         "define",
@@ -426,6 +428,7 @@ fn parses_function_with_keyword_name() {
         "fn",
         "function",
         "constant",
+        "error",
         "takes",
         "returns",
         "define",
@@ -475,6 +478,7 @@ fn parses_function_with_keyword_name_in_macro() {
         "fn",
         "function",
         "constant",
+        "error",
         "takes",
         "returns",
         "define",
@@ -574,6 +578,7 @@ fn parses_keyword_arbitrary_whitespace() {
         ("macro", TokenKind::Macro),
         ("fn", TokenKind::Fn),
         ("constant", TokenKind::Constant),
+        ("error", TokenKind::Error),
         ("function", TokenKind::Function),
     ];
 
diff --git a/huff_parser/tests/error.rs b/huff_parser/tests/error.rs
new file mode 100644
index 0000000..8ffcc6a
--- /dev/null
+++ b/huff_parser/tests/error.rs
@@ -0,0 +1,31 @@
+use huff_lexer::*;
+use huff_parser::*;
+use huff_utils::prelude::*;
+
+#[test]
+fn test_parses_custom_error() {
+    let source = "#define error TestError(uint256)";
+    let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(flattened_source);
+    let tokens = lexer.into_iter().map(|x| x.unwrap()).collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, None);
+    let contract = parser.parse().unwrap();
+    assert_eq!(parser.current_token.kind, TokenKind::Eof);
+
+    let custom_error = contract.errors[0].clone();
+    assert_eq!(
+        custom_error,
+        ErrorDefinition {
+            name: String::from("TestError"),
+            selector: [124, 104, 44, 83],
+            span: AstSpan(vec![
+                Span { start: 0, end: 7, file: None },
+                Span { start: 8, end: 13, file: None },
+                Span { start: 14, end: 23, file: None },
+                Span { start: 23, end: 24, file: None },
+                Span { start: 24, end: 31, file: None },
+                Span { start: 31, end: 32, file: None }
+            ])
+        }
+    );
+}