support #[hdl] on functions -- enables #[hdl] usage in function body

crates/fayalite-proc-macros-impl/src/module/transform_body.rs

@@ -925,7 +925,7 @@ with_debug_clone_and_fold! {
     #[allow(dead_code)]
     pub(crate) struct HdlLet<Kind = HdlLetKind> {
         pub(crate) attrs: Vec<Attribute>,
-        pub(crate) hdl_attr: HdlAttr<Nothing>,
+        pub(crate) hdl_attr: HdlAttr<Nothing, kw::hdl>,
         pub(crate) let_token: Token![let],
         pub(crate) mut_token: Option<Token![mut]>,
         pub(crate) name: Ident,
@@ -1112,7 +1112,7 @@ impl<T: ToString> ToTokens for ImplicitName<T> {
 }
 
 struct Visitor<'a> {
-    module_kind: ModuleKind,
+    module_kind: Option<ModuleKind>,
     errors: Errors,
     io: Vec<ModuleIO>,
     block_depth: usize,
@@ -1120,22 +1120,33 @@ struct Visitor<'a> {
 }
 
 impl Visitor<'_> {
-    fn take_hdl_attr<T: Parse>(&mut self, attrs: &mut Vec<Attribute>) -> Option<HdlAttr<T>> {
+    fn take_hdl_attr<T: Parse>(
+        &mut self,
+        attrs: &mut Vec<Attribute>,
+    ) -> Option<HdlAttr<T, kw::hdl>> {
         self.errors.unwrap_or(
             HdlAttr::parse_and_take_attr(attrs),
             Some(syn::parse2::<T>(quote! {}).unwrap().into()),
         )
     }
-    fn require_normal_module(&mut self, spanned: impl ToTokens) {
+    fn require_normal_module_or_fn(&mut self, spanned: impl ToTokens) {
         match self.module_kind {
-            ModuleKind::Extern => {
+            Some(ModuleKind::Extern) => {
                 self.errors
                     .error(spanned, "not allowed in #[hdl_module(extern)]");
             }
-            ModuleKind::Normal => {}
+            Some(ModuleKind::Normal) | None => {}
         }
     }
-    fn process_hdl_if(&mut self, hdl_attr: HdlAttr<Nothing>, expr_if: ExprIf) -> Expr {
+    fn require_module(&mut self, spanned: impl ToTokens) {
+        match self.module_kind {
+            None => {
+                self.errors.error(spanned, "not allowed in #[hdl] fn");
+            }
+            Some(_) => {}
+        }
+    }
+    fn process_hdl_if(&mut self, hdl_attr: HdlAttr<Nothing, kw::hdl>, expr_if: ExprIf) -> Expr {
         let ExprIf {
             attrs,
             if_token,
@@ -1143,7 +1154,7 @@ impl Visitor<'_> {
             then_branch,
             else_branch,
         } = expr_if;
-        self.require_normal_module(if_token);
+        self.require_normal_module_or_fn(if_token);
         let else_expr = else_branch.unzip().1.map(|else_expr| match *else_expr {
             Expr::If(expr_if) => self.process_hdl_if(hdl_attr.clone(), expr_if),
             expr => expr,
@@ -1208,11 +1219,12 @@ impl Visitor<'_> {
             .to_tokens(expr);
         });
         let mut attrs = hdl_let.attrs.clone();
+        self.require_module(kind);
         match self.module_kind {
-            ModuleKind::Extern => attrs.push(parse_quote_spanned! {hdl_let.let_token.span=>
+            Some(ModuleKind::Extern) => attrs.push(parse_quote_spanned! {hdl_let.let_token.span=>
                 #[allow(unused_variables)]
             }),
-            ModuleKind::Normal => {}
+            Some(ModuleKind::Normal) | None => {}
         }
         let let_stmt = Local {
             attrs,
@@ -1249,7 +1261,7 @@ impl Visitor<'_> {
                 },
             semi_token,
         } = hdl_let;
-        self.require_normal_module(instance);
+        self.require_normal_module_or_fn(instance);
         let mut expr = instance.to_token_stream();
         paren.surround(&mut expr, |expr| {
             let name_str = ImplicitName {
@@ -1276,7 +1288,7 @@ impl Visitor<'_> {
     fn process_hdl_let_reg_builder(&mut self, hdl_let: HdlLet<HdlLetKindRegBuilder>) -> Local {
         let name = &hdl_let.name;
         let reg_builder = hdl_let.kind.reg_builder;
-        self.require_normal_module(reg_builder);
+        self.require_normal_module_or_fn(reg_builder);
         let mut expr = reg_builder.to_token_stream();
         hdl_let.kind.reg_builder_paren.surround(&mut expr, |expr| {
             let name_str = ImplicitName {
@@ -1327,7 +1339,7 @@ impl Visitor<'_> {
     fn process_hdl_let_wire(&mut self, hdl_let: HdlLet<HdlLetKindWire>) -> Local {
         let name = &hdl_let.name;
         let wire = hdl_let.kind.wire;
-        self.require_normal_module(wire);
+        self.require_normal_module_or_fn(wire);
         let ty_expr = unwrap_or_static_type(hdl_let.kind.ty_expr.as_ref(), wire.span());
         let mut expr = wire.to_token_stream();
         hdl_let.kind.paren.surround(&mut expr, |expr| {
@@ -1361,7 +1373,7 @@ impl Visitor<'_> {
         let name = &hdl_let.name;
         let memory_fn = hdl_let.kind.memory_fn;
         let memory_fn_name = memory_fn.name();
-        self.require_normal_module(memory_fn_name);
+        self.require_normal_module_or_fn(memory_fn_name);
         let mut expr = memory_fn_name.to_token_stream();
         let (paren, arg) = match memory_fn {
             MemoryFn::Memory {
@@ -1543,7 +1555,7 @@ impl Fold for Visitor<'_> {
     }
 
     fn fold_attribute(&mut self, attr: Attribute) -> Attribute {
-        if is_hdl_attr(&attr) {
+        if is_hdl_attr::<kw::hdl>(&attr) {
             self.errors
                 .error(&attr, "#[hdl] attribute not supported here");
         }
@@ -1610,8 +1622,9 @@ impl Fold for Visitor<'_> {
     fn fold_local(&mut self, let_stmt: Local) -> Local {
         match self
             .errors
-            .ok(HdlAttr::<Nothing>::parse_and_leave_attr(&let_stmt.attrs))
-        {
+            .ok(HdlAttr::<Nothing, kw::hdl>::parse_and_leave_attr(
+                &let_stmt.attrs,
+            )) {
             None => return empty_let(),
             Some(None) => return fold_local(self, let_stmt),
             Some(Some(HdlAttr { .. })) => {}
@@ -1646,7 +1659,7 @@ impl Fold for Visitor<'_> {
 }
 
 pub(crate) fn transform_body(
-    module_kind: ModuleKind,
+    module_kind: Option<ModuleKind>,
     mut body: Box<Block>,
     parsed_generics: &ParsedGenerics,
 ) -> syn::Result<(Box<Block>, Vec<ModuleIO>)> {