Add optional parametrization for "reg" keyword

Author: VonTum

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "sus_compiler"
 description = "Compiler for the SUS Hardware Design Language"
-version = "0.4.0"
+version = "0.5.0-dev"
 authors = ["Lennart Van Hirtum <lennart.vanhirtum@gmail.com>"]
 license = "GPL-3.0-or-later"
 repository = "https://github.com/pc2/sus-compiler"
@@ -42,9 +42,9 @@ pulldown-cmark = { version = "0.13", default-features = false, features = ["html
 
 # Tree sitter
 tree-sitter = "0.26.6"
-tree-sitter-sus = {version = "~0.4.0", path = "./tree-sitter-sus"}
+tree-sitter-sus = {version = "~0.5.0-dev", path = "./tree-sitter-sus"}
 
-sus-proc-macro = {version = "~0.4.0", path = "./sus-proc-macro"}
+sus-proc-macro = {version = "~0.5.0-dev", path = "./sus-proc-macro"}
 
 # calyx-ir = {version = "0.6.1", optional = true}
 # calyx-opt = {version = "0.6.1", optional = true}

src/flattening/flatten.rs

@@ -1690,36 +1690,61 @@ impl<'l, 'c: 'l> FlatteningContext<'l, '_> {
         self.alloc_local_name(clock_name_span, clock_name, NamedLocal::ClockDecl(clock_id));
     }
 
-    fn flatten_write_modifiers(&self, cursor: &mut Cursor<'c>) -> WriteModifiers {
+    fn flatten_reg_modifier(&mut self, cursor: &mut Cursor<'c>) -> RegExpression {
+        let whole_span = cursor.span();
+        cursor.go_down(kind!("reg_modifier"), move |cursor| {
+            let reg_parameter = if cursor.optional_field(field!("reg_param")) {
+                let bracket_span = BracketSpan::from_outer(cursor.span());
+                cursor.go_down_content(kind!("parenthesis_expression"), move |cursor| {
+                    let expr_id = self.flatten_subexpr(cursor);
+                    Some((expr_id, bracket_span))
+                })
+            } else {
+                None
+            };
+            RegExpression {
+                whole_span,
+                reg_parameter,
+            }
+        })
+    }
+
+    fn flatten_write_modifiers(&mut self, cursor: &mut Cursor<'c>) -> WriteModifiers {
         if cursor.optional_field(field!("write_modifiers")) {
-            let modifiers_span = cursor.span();
-            let mut initial_count = 0;
-            let mut reg_count = 0;
+            let mut initial_kw_span: Option<Span> = None;
+            let mut regs: Vec<RegExpression> = Vec::new();
             cursor.list(kind!("write_modifiers"), |cursor| {
                 let kw_kind = cursor.kind();
-                if kw_kind == kw!("reg") {
-                    reg_count += 1;
+                if kw_kind == kind!("reg_modifier") {
+                    regs.push(self.flatten_reg_modifier(cursor));
                 } else if kw_kind == kw!("initial") {
-                    initial_count += 1;
+                    let span = cursor.span();
+                    if let Some(prev_kw) = initial_kw_span {
+                        self.errors
+                            .error(span, "This write already has an 'initial' keyword")
+                            .info(prev_kw, "Previous 'initial' keyword");
+                    } else {
+                        initial_kw_span = Some(span);
+                    }
                 } else {
                     unreachable!()
                 }
             });
-            match (initial_count, reg_count) {
-                (0, num_regs) => WriteModifiers::Connection {
-                    num_regs,
-                    regs_span: modifiers_span,
-                },
-                (1, 0) => WriteModifiers::Initial {
-                    initial_kw_span: modifiers_span,
-                },
-                _other => unreachable!(),
+            if let Some(initial_kw_span) = initial_kw_span {
+                for r in &regs {
+                    self.errors
+                        .error(
+                            r.whole_span,
+                            "Cannot both place pipeline regs on an 'initial' write.",
+                        )
+                        .info(initial_kw_span, "'Initial' keyword used here");
+                }
+                WriteModifiers::Initial { initial_kw_span }
+            } else {
+                WriteModifiers::Connection { regs }
             }
         } else {
-            WriteModifiers::Connection {
-                num_regs: 0,
-                regs_span: cursor.span().empty_span_at_front(),
-            }
+            WriteModifiers::Connection { regs: Vec::new() }
         }
     }
 

src/flattening/mod.rs

@@ -401,6 +401,12 @@ impl WireReference {
     }
 }
 
+#[derive(Debug)]
+pub struct RegExpression {
+    pub whole_span: Span,
+    pub reg_parameter: Option<(FlatID, BracketSpan)>,
+}
+
 /// In a [Write], this represents what kind of write it is, based on keywords `reg` or `initial`
 #[derive(Debug)]
 pub enum WriteModifiers {
@@ -409,7 +415,7 @@ pub enum WriteModifiers {
     /// int v
     /// reg reg v = a * 3
     /// ```
-    Connection { num_regs: i64, regs_span: Span },
+    Connection { regs: Vec<RegExpression> },
     /// Set the initial value of a `state` register
     /// ```sus
     /// state int count

src/flattening/typecheck/clock_check.rs

@@ -94,8 +94,8 @@ impl<'l> TypeCheckingContext<'l> {
                             .unwrap_or(Physical(ClockID::UNKNOWN));
                         let mut target_span = wr.to.root_span;
 
-                        match wr.write_modifiers {
-                            WriteModifiers::Connection { .. } => {
+                        match &wr.write_modifiers {
+                            WriteModifiers::Connection { regs } => {
                                 if let Some(condition_domain) =
                                     self.get_condition_domain(expr.parent_condition)
                                     && let Physical(target_phys) = &mut target_domain
@@ -108,10 +108,18 @@ impl<'l> TypeCheckingContext<'l> {
                                         "the runtime condition",
                                     );
                                 }
+                                for r in regs {
+                                    if let Some(param) = r.reg_parameter {
+                                        self.must_be_generative(
+                                            param.0,
+                                            "the number of pipeline regs parameter",
+                                        );
+                                    }
+                                }
                             }
                             WriteModifiers::Initial { initial_kw_span } => {
                                 target_domain = Generative;
-                                target_span = initial_kw_span;
+                                target_span = *initial_kw_span;
                             }
                         }
                         wr.target_domain.set_initial(target_domain);

src/instantiation/execute.rs

@@ -863,16 +863,25 @@ impl<'l> ExecutionContext<'l> {
         domain: ClockID,
     ) -> ExecutionResult<()> {
         let_unwrap!(
-            WriteModifiers::Connection {
-                num_regs,
-                regs_span: _,
-            },
+            WriteModifiers::Connection { regs },
             &write_to.write_modifiers
         );
+
+        let mut num_regs = 0;
+        for r in regs {
+            if let Some(param) = &r.reg_parameter {
+                num_regs += self
+                    .generation_state
+                    .get_generation_small_int::<i64>(param.0)?;
+            } else {
+                num_regs += 1;
+            }
+        }
+
         let (target_wire, path) =
             self.wire_ref_to_real_path(&write_to.to, original_instruction, domain)?;
 
-        self.instantiate_write_to_wire(target_wire, path, from, *num_regs, write_span);
+        self.instantiate_write_to_wire(target_wire, path, from, num_regs, write_span);
         Ok(())
     }
 

src/to_string.rs

@@ -427,12 +427,17 @@ impl UniCell<ClockDomain> {
 impl Display for WriteModifiers {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
-            WriteModifiers::Connection { num_regs, .. } => {
-                for _ in 0..*num_regs {
-                    f.write_str("reg ")?;
-                }
-                Ok(())
-            }
+            WriteModifiers::Connection { regs } => write!(
+                f,
+                "{}",
+                display_join(" ", regs.iter(), |f, r| {
+                    f.write_str("reg")?;
+                    if r.reg_parameter.is_some() {
+                        f.write_str("(...)")?;
+                    }
+                    Ok(())
+                })
+            ),
             WriteModifiers::Initial { .. } => f.write_str("initial"),
         }
     }

sus-proc-macro/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "sus-proc-macro"
 description = "Proc-macros for use in the sus-compiler"
-version = "0.4.0"
+version = "0.5.0-dev"
 authors = ["Lennart Van Hirtum <lennart.vanhirtum@gmail.com>"]
 license = "GPL-3.0-or-later"
 repository = "https://github.com/pc2/sus-compiler"
@@ -15,7 +15,7 @@ edition = "2024"
 
 # Tree sitter
 tree-sitter = "0.26.6"
-tree-sitter-sus = {version = "~0.4.0", path = "../tree-sitter-sus"}
+tree-sitter-sus = {version = "~0.5.0-dev", path = "../tree-sitter-sus"}
 regex = "1.11.1"
 
 # proc_macro utils

test.sus

@@ -1888,3 +1888,8 @@ module PositiveAndNegativeLatencyOffset {
 	LatencyOffset#(T: type bool, OFFSET: 3) pos
 	LatencyOffset#(T: type bool, OFFSET: -3) neg
 }
+
+module ParametrizeableRegs {
+	int x = 5
+	reg(5) int y = x
+}