2 files changed, 403 insertions, 225 deletions

diff --git a/src/Language/Fiddle/Ast/Internal/Stage.hs b/src/Language/Fiddle/Ast/Internal/Stage.hs
index cfcd0e6..20460b6 100644
--- a/src/Language/Fiddle/Ast/Internal/Stage.hs
+++ b/src/Language/Fiddle/Ast/Internal/Stage.hs
@@ -3,6 +3,7 @@
 {-# LANGUAGE IncoherentInstances #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
+
 module Language.Fiddle.Ast.Internal.Stage where
 
 import Data.Type.Bool
@@ -11,33 +12,43 @@ import Data.Typeable
 import qualified GHC.TypeError as TypeError
 import GHC.TypeLits
 
--- Stage of compilation. Parts of the AST maybe un unavailable with other stages
--- as compilation simplifies the AST.
-data Stage = Stage1 | Stage2 | Stage3
+-- | Represents the different stages of the compilation process.
+--   Different stages may have distinct versions of the Abstract Syntax Tree (AST),
+--   as the compilation process simplifies or transforms the tree.
+data Stage
+  = Parsed
+  | Expanded
+  | Checked
   deriving (Typeable)
 
--- Returns the stage before the given stage.
-type family PreviousStage (s :: Stage) :: Stage where
-  PreviousStage Stage1 = TypeError (TypeError.Text "No Prior stage to Stage1")
-  PreviousStage s = NumberToStage (StageToNumber s - 1)
+-- | Converts a 'Stage' into a type-level natural number. This mapping allows
+--   us to use numeric comparisons at the type level to reason about stages.
+--   For example, 'Parsed' maps to 1, 'Expanded' to 2, and so on.
+type family StageToNumber (s :: Stage) :: Natural where
+  StageToNumber Parsed = 1
+  StageToNumber Expanded = 2
+  StageToNumber Checked = 3
 
--- Returns the stage after the give stage.
-type family NextStage (s :: Stage) :: Stage where
-  NextStage s = NumberToStage (StageToNumber s + 1)
+-- | A type-level constraint that checks if one compilation stage precedes another.
+--   It compares the numeric values associated with each stage using 'CmpNat'.
+--   If 'stage1' is less than 'stage2', the result will be 'True'.
+--   This is useful to conditionally include or exclude parts of the AST
+--   depending on the compilation stage.
+type StagePreceeds stage1 stage2 =
+  (CmpNat (StageToNumber stage1) (StageToNumber stage2) == LT)
 
--- Converts a stage to a number.
-type family StageToNumber (s :: Stage) :: Natural where
-  StageToNumber Stage1 = 1
-  StageToNumber Stage2 = 2
-  StageToNumber Stage3 = 3
-
--- Converts a number to a stage.
-type family NumberToStage (n :: Natural) :: Stage where
-  NumberToStage 1 = Stage1
-  NumberToStage 2 = Stage2
-  NumberToStage 3 = Stage3
-  NumberToStage n = TypeError (TypeError.Text "Number in NumberToStage.")
-
--- Type-level constraint to determine if a stage is less than some natural
--- ordinal. Used to bound parts of the AST in multiple stages.
-type StageLessThan stage (n :: Natural) = (CmpNat (StageToNumber stage) n == LT)
+type (<) a b = StagePreceeds a b
+
+-- | A type-level constraint that checks if one compilation stage succeeds another.
+--   Similar to 'StagePreceeds', it compares the numeric values of stages.
+--   Returns 'True' if 'stage1' comes after 'stage2'.
+type StageSucceeds stage1 stage2 =
+  (CmpNat (StageToNumber stage1) (StageToNumber stage2) == LT)
+
+type (>) a b = StagePreceeds a b
+
+-- | A type-level function that compares two stages and returns a comparison
+--   result ('LT', 'EQ', or 'GT'). This function is a generalized way to compare
+--   the order of stages numerically.
+type CmpStage stage1 stage2 =
+  CmpNat (StageToNumber stage1) (StageToNumber stage2)
diff --git a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
index 827f712..d03a855 100644
--- a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
+++ b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
@@ -73,18 +73,16 @@ import Language.Fiddle.Ast.Internal.Stage
 -- numbers are just strings like anything else. In later stages, numbers get
 -- parsed into actual integers. This makes it easier to process later.
 type family NumberType (a :: Stage) :: Type where
-  NumberType Stage1 = Text
-  NumberType Stage2 = Integer
-  NumberType s = NumberType (PreviousStage s)
+  NumberType s = If (s < Expanded) Text Integer
 
 -- The type that represents an import statement. In the early stages of
 -- compilation, this is just a string representing the import path, but in later
 -- stages of compilation, this actually gets replaced by an abstract
 -- representation of the imported material.
 type family ImportType (stage :: Stage) :: SynTree where
-  ImportType Stage1 = ImportStatement
-  ImportType Stage2 = ImportStatement
-  ImportType Stage3 = ImportStatement
+  ImportType Parsed = ImportStatement
+  ImportType Expanded = ImportStatement
+  ImportType Checked = ImportStatement
 
 -- A way to disable or enable a subtree type based on a type-level boolean.
 --
@@ -109,66 +107,91 @@ data Name :: SynTree where
   Name :: NonEmpty (Identifier f a) -> a -> Name f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- | Represents a directive in the Fiddle language. A directive provides 
---   additional metadata or instructions that the compiler can use during 
---   code generation. Directives can be attached to many elements in the 
+-- | Represents a directive in the Fiddle language. A directive provides
+--   additional metadata or instructions that the compiler can use during
+--   code generation. Directives can be attached to many elements in the
 --   syntax tree.
 data Directive :: SynTree where
   Directive ::
-    { directiveBody :: f (DirectiveBody f a), -- ^ The body of the directive.
-      directiveAnnot :: a -- ^ Annotation for the directive.
-    } -> Directive f a
+    { -- | The body of the directive.
+      directiveBody :: f (DirectiveBody f a),
+      -- | Annotation for the directive.
+      directiveAnnot :: a
+    } ->
+    Directive f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents the body of a directive, which consists of multiple elements.
 data DirectiveBody :: SynTree where
   DirectiveBody ::
-    { directiveElements :: [DirectiveElement f a], -- ^ Elements of the directive.
-      directiveBodyAnnot :: a -- ^ Annotation for the directive body.
-    } -> DirectiveBody f a
+    { -- | Elements of the directive.
+      directiveElements :: [DirectiveElement f a],
+      -- | Annotation for the directive body.
+      directiveBodyAnnot :: a
+    } ->
+    DirectiveBody f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- | Represents an element in a directive. Can be either a key or a key-value 
+-- | Represents an element in a directive. Can be either a key or a key-value
 --   pair.
 data DirectiveElement :: SynTree where
   -- | A simple directive element with a key. The mere presence of this key
   --   holds semantic value.
   DirectiveElementKey ::
-    { directiveBackend :: Maybe (Identifier f a), -- ^ Optional backend target.
-      directiveKey :: Identifier f a, -- ^ The key of the directive.
-      directiveKeyAnnot :: a -- ^ Annotation for the directive element.
-    } -> DirectiveElement f a
+    { -- | Optional backend target.
+      directiveBackend :: Maybe (Identifier f a),
+      -- | The key of the directive.
+      directiveKey :: Identifier f a,
+      -- | Annotation for the directive element.
+      directiveKeyAnnot :: a
+    } ->
+    DirectiveElement f a
   -- | A more complex directive element with a key-value pair, optionally
   --   specifying a backend.
   DirectiveElementKeyValue ::
-    { directiveBackend :: Maybe (Identifier f a), -- ^ Optional backend target.
-      directiveKey :: Identifier f a, -- ^ The key of the directive.
-      directiveValue :: DirectiveExpression f a, -- ^ The value of the directive.
-      directiveKeyValueAnnot :: a -- ^ Annotation for the key-value directive.
-    } -> DirectiveElement f a
+    { -- | Optional backend target.
+      directiveBackend :: Maybe (Identifier f a),
+      -- | The key of the directive.
+      directiveKey :: Identifier f a,
+      -- | The value of the directive.
+      directiveValue :: DirectiveExpression f a,
+      -- | Annotation for the key-value directive.
+      directiveKeyValueAnnot :: a
+    } ->
+    DirectiveElement f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- | Represents expressions that can be used within a directive, either a 
+-- | Represents expressions that can be used within a directive, either a
 --   string or a number.
 data DirectiveExpression f a where
-  DirectiveString :: 
-    { directiveStringValue :: Text, -- ^ String value of the directive.
-      directiveStringAnnot :: a -- ^ Annotation for the directive string.
-    } -> DirectiveExpression f a
-  DirectiveNumber :: 
-    { directiveNumberValue :: Text, -- ^ Number value of the directive.
-      directiveNumberAnnot :: a -- ^ Annotation for the directive number.
-    } -> DirectiveExpression f a
+  DirectiveString ::
+    { -- | String value of the directive.
+      directiveStringValue :: Text,
+      -- | Annotation for the directive string.
+      directiveStringAnnot :: a
+    } ->
+    DirectiveExpression f a
+  DirectiveNumber ::
+    { -- | Number value of the directive.
+      directiveNumberValue :: Text,
+      -- | Annotation for the directive number.
+      directiveNumberAnnot :: a
+    } ->
+    DirectiveExpression f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- | A type that wraps another syntax tree and applies a list of directives to 
+-- | A type that wraps another syntax tree and applies a list of directives to
 --   it.
 data Directed t stage f a where
   Directed ::
-    { directedDirectives :: [Directive f a], -- ^ List of directives.
-      directedSubtree :: t stage f a, -- ^ The wrapped syntax tree.
-      directedAnnot :: a -- ^ Annotation for the directed subtree.
-    } -> Directed t stage f a
+    { -- | List of directives.
+      directedDirectives :: [Directive f a],
+      -- | The wrapped syntax tree.
+      directedSubtree :: t stage f a,
+      -- | Annotation for the directed subtree.
+      directedAnnot :: a
+    } ->
+    Directed t stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Apply a function to the underlying subtree in a 'Directed' type.
@@ -177,119 +200,168 @@ mapDirected fn (Directed dr tfa a) = Directed dr (fn tfa) a
 
 -- | Apply a monadic function to the underlying subtree in a 'Directed' type.
 mapDirectedM ::
-  (Monad m) => 
-  (t s f a -> m (t' s' f a)) -> 
-  Directed t s f a -> 
+  (Monad m) =>
+  (t s f a -> m (t' s' f a)) ->
+  Directed t s f a ->
   m (Directed t' s' f a)
 mapDirectedM fn (Directed dr tfa a) = Directed dr <$> fn tfa <*> pure a
 
--- | Convert an annotated syntax tree element to a 'Directed' type with 
+-- | Convert an annotated syntax tree element to a 'Directed' type with
 --   an empty directive list.
 asDirected :: (Annotated (t s)) => t s f a -> Directed t s f a
 asDirected tfa = Directed [] tfa (annot tfa)
 
--- | Extract the underlying subtree from a 'Directed' type, discarding any 
+-- | Extract the underlying subtree from a 'Directed' type, discarding any
 --   directives.
 undirected :: Directed t s f a -> t s f a
 undirected (Directed _ tfa _) = tfa
 
 -- | The root of the parse tree, containing a list of top-level declarations.
 data FiddleUnit (stage :: Stage) (f :: Type -> Type) a where
-  FiddleUnit :: 
-    { fiddleDecls :: [Directed FiddleDecl stage f a], -- ^ List of declarations.
-      fiddleUnitAnnot :: a -- ^ Annotation for the 'FiddleUnit'.
-    } -> FiddleUnit stage f a
+  FiddleUnit ::
+    { -- | List of declarations.
+      fiddleDecls :: [Directed FiddleDecl stage f a],
+      -- | Annotation for the 'FiddleUnit'.
+      fiddleUnitAnnot :: a
+    } ->
+    FiddleUnit stage f a
   deriving (Generic, Annotated, Typeable)
 
 deriving instance (Alter (ImportType stage)) => Alter (FiddleUnit stage)
 
 -- | Represents an identifier with an associated annotation.
-data Identifier f a = Identifier 
-  { identifierName :: !Text, -- ^ The name of the identifier.
-    identifierAnnot :: a -- ^ Annotation for the identifier.
+data Identifier f a = Identifier
+  { -- | The name of the identifier.
+    identifierName :: !Text,
+    -- | Annotation for the identifier.
+    identifierAnnot :: a
   }
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Expressions used within Fiddle, including literals and variables.
 data Expression (s :: Stage) :: SynTree where
   -- | A numeric literal, whose value is dependent on the compilation stage.
-  LitNum :: 
-    { litNumValue :: NumberType stage, -- ^ The numeric value.
-      litNumAnnot :: a -- ^ Annotation for the literal.
-    } -> Expression stage f a
+  LitNum ::
+    { -- | The numeric value.
+      litNumValue :: NumberType stage,
+      -- | Annotation for the literal.
+      litNumAnnot :: a
+    } ->
+    Expression stage f a
   -- | A variable reference.
-  Var :: 
-    { varIdentifier :: Identifier f a, -- ^ The identifier of the variable.
-      varAnnot :: a -- ^ Annotation for the variable.
-    } -> Expression stage f a
+  Var ::
+    { -- | The identifier of the variable.
+      varIdentifier :: Identifier f a,
+      -- | Annotation for the variable.
+      varAnnot :: a
+    } ->
+    Expression stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents an import statement in the Fiddle language.
 data ImportStatement f a where
-  ImportStatement :: 
-    { importPath :: Text, -- ^ The path to import.
-      importList :: Maybe (ImportList f a), -- ^ Optional list of imported items.
-      importStatementAnnot :: a -- ^ Annotation for the import statement.
-    } -> ImportStatement f a
+  ImportStatement ::
+    { -- | The path to import.
+      importPath :: Text,
+      -- | Optional list of imported items.
+      importList :: Maybe (ImportList f a),
+      -- | Annotation for the import statement.
+      importStatementAnnot :: a
+    } ->
+    ImportStatement f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | A list of imported identifiers.
 data ImportList f a where
-  ImportList :: 
-    { importIdentifiers :: [Identifier f a], -- ^ The list of identifiers.
-      importListAnnot :: a -- ^ Annotation for the import list.
-    } -> ImportList f a
+  ImportList ::
+    { -- | The list of identifiers.
+      importIdentifiers :: [Identifier f a],
+      -- | Annotation for the import list.
+      importListAnnot :: a
+    } ->
+    ImportList f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents top-level declarations in Fiddle.
 data FiddleDecl :: StagedSynTree where
   -- | An option declaration in the form 'option <ident> <ident>'.
   OptionDecl ::
-    { optionKey :: Identifier f a, -- ^ The key of the option.
-      optionValue :: Identifier f a, -- ^ The value of the option.
-      optionAnnot :: a -- ^ Annotation for the option declaration.
-    } -> FiddleDecl stage f a
+    { -- | The key of the option.
+      optionKey :: Identifier f a,
+      -- | The value of the option.
+      optionValue :: Identifier f a,
+      -- | Annotation for the option declaration.
+      optionAnnot :: a
+    } ->
+    FiddleDecl stage f a
   -- | An import declaration.
   ImportDecl ::
-    { importType :: ImportType stage f a, -- ^ The imported type.
-      importDeclAnnot :: a -- ^ Annotation for the import declaration.
-    } -> FiddleDecl stage f a
+    { -- | The imported type.
+      importType :: ImportType stage f a,
+      -- | Annotation for the import declaration.
+      importDeclAnnot :: a
+    } ->
+    FiddleDecl stage f a
   -- | A using declaration.
   UsingDecl ::
-    { usingName :: Name f a, -- ^ The name being used.
-      usingAnnot :: a -- ^ Annotation for the using declaration.
-    } -> FiddleDecl stage f a
+    { -- | The name being used.
+      usingName :: Name f a,
+      -- | Annotation for the using declaration.
+      usingAnnot :: a
+    } ->
+    FiddleDecl stage f a
   -- | A package declaration.
   PackageDecl ::
-    { packageName :: Name f a, -- ^ The package name.
-      packageBody :: f (PackageBody stage f a), -- ^ The body of the package.
-      packageAnnot :: a -- ^ Annotation for the package declaration.
-    } -> FiddleDecl stage f a
+    { -- | The package name.
+      packageName :: Name f a,
+      -- | The body of the package.
+      packageBody :: f (PackageBody stage f a),
+      -- | Annotation for the package declaration.
+      packageAnnot :: a
+    } ->
+    FiddleDecl stage f a
   -- | A location declaration in the form 'location <identifier> = <expr>'.
   LocationDecl ::
-    { locationIdent :: Identifier f a, -- ^ The location identifier.
-      locationExpr :: Expression stage f a, -- ^ The associated expression.
-      locationAnnot :: a -- ^ Annotation for the location declaration.
-    } -> FiddleDecl stage f a
+    { -- | The location identifier.
+      locationIdent :: Identifier f a,
+      -- | The associated expression.
+      locationExpr :: Expression stage f a,
+      -- | Annotation for the location declaration.
+      locationAnnot :: a
+    } ->
+    FiddleDecl stage f a
   -- | A bits declaration in the form 'bits <identifier> : <type>'.
   BitsDecl ::
-    { bitsIdent :: Identifier f a, -- ^ The identifier of the bits.
-      bitsType :: BitType stage f a, -- ^ The type of the bits.
-      bitsAnnot :: a -- ^ Annotation for the bits declaration.
-    } -> FiddleDecl stage f a
+    { -- | The identifier of the bits.
+      bitsIdent :: Identifier f a,
+      -- | The type of the bits.
+      bitsType :: BitType stage f a,
+      -- | Annotation for the bits declaration.
+      bitsAnnot :: a
+    } ->
+    FiddleDecl stage f a
   -- | An object type declaration.
   ObjTypeDecl ::
-    { objTypeIdent :: Identifier f a, -- ^ The identifier of the object type.
-      objTypeBody :: f (ObjTypeBody stage f a), -- ^ The body of the object type.
-      objTypeAnnot :: a -- ^ Annotation for the object type declaration.
-    } -> FiddleDecl stage f a
+    { -- | The identifier of the object type.
+      objTypeIdent :: Identifier f a,
+      -- | The body of the object type.
+      objTypeBody :: f (ObjTypeBody stage f a),
+      -- | Annotation for the object type declaration.
+      objTypeAnnot :: a
+    } ->
+    FiddleDecl stage f a
   -- | An object declaration in the form 'object <ident> at <expr> : <type>'.
   ObjectDecl ::
-    { objectIdent :: Identifier f a, -- ^ The identifier of the object.
-      objectLocation :: Expression stage f a, -- ^ The location expression.
-      objectType :: ObjType stage f a, -- ^ The type of the object.
-      objectAnnot :: a -- ^ Annotation for the object declaration.
-    } -> FiddleDecl stage f a
+    { -- | The identifier of the object.
+      objectIdent :: Identifier f a,
+      -- | The location expression.
+      objectLocation :: Expression stage f a,
+      -- | The type of the object.
+      objectType :: ObjType stage f a,
+      -- | Annotation for the object declaration.
+      objectAnnot :: a
+    } ->
+    FiddleDecl stage f a
   deriving (Generic, Annotated, Typeable)
 
 deriving instance (Alter (ImportType stage)) => Alter (FiddleDecl stage)
@@ -298,32 +370,47 @@ deriving instance (Alter (ImportType stage)) => Alter (FiddleDecl stage)
 --   union), a list of object declarations, and an annotation.
 data ObjTypeBody (stage :: Stage) (f :: Type -> Type) a where
   ObjTypeBody ::
-    { objBodyType :: BodyType f a, -- ^ The body type (struct or union).
-      objBodyDecls :: [Directed ObjTypeDecl stage f a], -- ^ Object declarations.
-      objBodyAnnot :: a -- ^ Annotation for the object type body.
-    } -> ObjTypeBody stage f a
+    { -- | The body type (struct or union).
+      objBodyType :: BodyType f a,
+      -- | Object declarations.
+      objBodyDecls :: [Directed ObjTypeDecl stage f a],
+      -- | Annotation for the object type body.
+      objBodyAnnot :: a
+    } ->
+    ObjTypeBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents an object type, which can be anonymous, an array, or a
 --   reference to another type.
 data ObjType stage f a where
-  -- | An anonymous object type, allowed only in Stage1. 
+  -- | An anonymous object type, allowed only in Parsed.
   AnonymousObjType ::
-    { anonWitness :: Witness (stage == Stage1), -- ^ Witness for stage constraint.
-      anonBody :: f (ObjTypeBody stage f a), -- ^ The body of the anonymous type.
-      anonAnnot :: a -- ^ Annotation for the anonymous type.
-    } -> ObjType stage f a
+    { -- | Witness for stage constraint.
+      anonWitness :: Witness (stage == Parsed),
+      -- | The body of the anonymous type.
+      anonBody :: f (ObjTypeBody stage f a),
+      -- | Annotation for the anonymous type.
+      anonAnnot :: a
+    } ->
+    ObjType stage f a
   -- | An array of object types.
   ArrayObjType ::
-    { arrayObjType :: ObjType stage f a, -- ^ The type of the array elements.
-      arraySize :: Expression stage f a, -- ^ The size of the array.
-      arrayAnnot :: a -- ^ Annotation for the array type.
-    } -> ObjType stage f a
+    { -- | The type of the array elements.
+      arrayObjType :: ObjType stage f a,
+      -- | The size of the array.
+      arraySize :: Expression stage f a,
+      -- | Annotation for the array type.
+      arrayAnnot :: a
+    } ->
+    ObjType stage f a
   -- | A reference to an existing type by name.
   ReferencedObjType ::
-    { refName :: Name f a, -- ^ The name of the referenced type.
-      refAnnot :: a -- ^ Annotation for the referenced type.
-    } -> ObjType stage f a
+    { -- | The name of the referenced type.
+      refName :: Name f a,
+      -- | Annotation for the referenced type.
+      refAnnot :: a
+    } ->
+    ObjType stage f a
   deriving (Typeable, Generic, Alter, Annotated, Typeable)
 
 -- | Represents a declaration inside an object type, such as a register, an
@@ -331,37 +418,57 @@ data ObjType stage f a where
 data ObjTypeDecl stage f a where
   -- | An assertion statement for a specific position.
   AssertPosStatement ::
-    { assertWitness :: Witness (StageLessThan stage 3), -- ^ Witness for stage constraint.
-      assertExpr :: Expression stage f a, -- ^ The expression for the assertion.
-      assertAnnot :: a -- ^ Annotation for the assertion.
-    } -> ObjTypeDecl stage f a
+    { -- | Witness for stage constraint.
+      assertWitness :: Witness (stage < Checked),
+      -- | The expression for the assertion.
+      assertExpr :: Expression stage f a,
+      -- | Annotation for the assertion.
+      assertAnnot :: a
+    } ->
+    ObjTypeDecl stage f a
   -- | A register declaration.
   RegisterDecl ::
-    { regModifier :: Maybe (Modifier f a), -- ^ Optional register modifier.
-      regIdent :: Maybe (Identifier f a), -- ^ Optional register identifier.
-      regSize :: Expression stage f a, -- ^ Size of the register.
-      regBody :: Maybe (RegisterBody stage f a), -- ^ Optional register body.
-      regAnnot :: a -- ^ Annotation for the register declaration.
-    } -> ObjTypeDecl stage f a
+    { -- | Optional register modifier.
+      regModifier :: Maybe (Modifier f a),
+      -- | Optional register identifier.
+      regIdent :: Maybe (Identifier f a),
+      -- | Size of the register.
+      regSize :: Expression stage f a,
+      -- | Optional register body.
+      regBody :: Maybe (RegisterBody stage f a),
+      -- | Annotation for the register declaration.
+      regAnnot :: a
+    } ->
+    ObjTypeDecl stage f a
   -- | A reserved declaration for padding or alignment.
   ReservedDecl ::
-    { reservedExpr :: Expression stage f a, -- ^ The expression for reserved space.
-      reservedAnnot :: a -- ^ Annotation for the reserved declaration.
-    } -> ObjTypeDecl stage f a
+    { -- | The expression for reserved space.
+      reservedExpr :: Expression stage f a,
+      -- | Annotation for the reserved declaration.
+      reservedAnnot :: a
+    } ->
+    ObjTypeDecl stage f a
   -- | A declaration for a substructure (struct or union).
   TypeSubStructure ::
-    { subStructureBody :: f (ObjTypeBody stage f a), -- ^ The body of the substructure.
-      subStructureName :: Maybe (Identifier f a), -- ^ Optional name for the substructure.
-      subStructureAnnot :: a -- ^ Annotation for the substructure.
-    } -> ObjTypeDecl stage f a
+    { -- | The body of the substructure.
+      subStructureBody :: f (ObjTypeBody stage f a),
+      -- | Optional name for the substructure.
+      subStructureName :: Maybe (Identifier f a),
+      -- | Annotation for the substructure.
+      subStructureAnnot :: a
+    } ->
+    ObjTypeDecl stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents a modifier for registers (e.g., read-only, read-write).
 data Modifier f a where
   ModifierKeyword ::
-    { modifierKey :: ModifierKeyword, -- ^ The keyword for the modifier.
-      modifierAnnot :: a -- ^ Annotation for the modifier.
-    } -> Modifier f a
+    { -- | The keyword for the modifier.
+      modifierKey :: ModifierKeyword,
+      -- | Annotation for the modifier.
+      modifierAnnot :: a
+    } ->
+    Modifier f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Enumerates the different types of register modifiers.
@@ -372,28 +479,39 @@ data ModifierKeyword = Rw | Ro | Wo
 --   declarations.
 data DeferredRegisterBody stage f a where
   DeferredRegisterBody ::
-    { deferredBits :: [Directed RegisterBitsDecl stage f a], -- ^ Bit declarations.
-      deferredAnnot :: a -- ^ Annotation for the deferred register body.
-    } -> DeferredRegisterBody stage f a
+    { -- | Bit declarations.
+      deferredBits :: [Directed RegisterBitsDecl stage f a],
+      -- | Annotation for the deferred register body.
+      deferredAnnot :: a
+    } ->
+    DeferredRegisterBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents the body type (struct or union) in an object.
 data BodyType (f :: Type -> Type) a where
   Union ::
-    { unionAnnot :: a -- ^ Annotation for the union.
-    } -> BodyType f a
+    { -- | Annotation for the union.
+      unionAnnot :: a
+    } ->
+    BodyType f a
   Struct ::
-    { structAnnot :: a -- ^ Annotation for the struct.
-    } -> BodyType f a
+    { -- | Annotation for the struct.
+      structAnnot :: a
+    } ->
+    BodyType f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents a register body with a body type and deferred bit declarations.
 data RegisterBody stage f a where
   RegisterBody ::
-    { regBodyType :: BodyType f a, -- ^ The body type of the register.
-      regDeferredBody :: f (DeferredRegisterBody stage f a), -- ^ Deferred body.
-      regBodyAnnot :: a -- ^ Annotation for the register body.
-    } -> RegisterBody stage f a
+    { -- | The body type of the register.
+      regBodyType :: BodyType f a,
+      -- | Deferred body.
+      regDeferredBody :: f (DeferredRegisterBody stage f a),
+      -- | Annotation for the register body.
+      regBodyAnnot :: a
+    } ->
+    RegisterBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents declarations within a register, such as defined bits,
@@ -401,22 +519,34 @@ data RegisterBody stage f a where
 data RegisterBitsDecl stage f a where
   -- | Declaration for reserved bits.
   ReservedBits ::
-    { reservedBitsExpr :: Expression stage f a, -- ^ Expression for reserved bits.
-      reservedBitsAnnot :: a -- ^ Annotation for the reserved bits.
-    } -> RegisterBitsDecl stage f a
+    { -- | Expression for reserved bits.
+      reservedBitsExpr :: Expression stage f a,
+      -- | Annotation for the reserved bits.
+      reservedBitsAnnot :: a
+    } ->
+    RegisterBitsDecl stage f a
   -- | Declaration for defined bits in a register.
   DefinedBits ::
-    { definedBitsModifier :: Maybe (Modifier f a), -- ^ Optional modifier for the bits.
-      definedBitsIdent :: Identifier f a, -- ^ Identifier for the bits.
-      definedBitsTypeRef :: RegisterBitsTypeRef stage f a, -- ^ Type reference for the bits.
-      definedBitsAnnot :: a -- ^ Annotation for the defined bits.
-    } -> RegisterBitsDecl stage f a
+    { -- | Optional modifier for the bits.
+      definedBitsModifier :: Maybe (Modifier f a),
+      -- | Identifier for the bits.
+      definedBitsIdent :: Identifier f a,
+      -- | Type reference for the bits.
+      definedBitsTypeRef :: RegisterBitsTypeRef stage f a,
+      -- | Annotation for the defined bits.
+      definedBitsAnnot :: a
+    } ->
+    RegisterBitsDecl stage f a
   -- | Substructure within a register.
   BitsSubStructure ::
-    { bitsSubRegBody :: RegisterBody stage f a, -- ^ The body of the substructure.
-      bitsSubName :: Maybe (Identifier f a), -- ^ Optional name for the substructure.
-      bitsSubAnnot :: a -- ^ Annotation for the substructure.
-    } -> RegisterBitsDecl stage f a
+    { -- | The body of the substructure.
+      bitsSubRegBody :: RegisterBody stage f a,
+      -- | Optional name for the substructure.
+      bitsSubName :: Maybe (Identifier f a),
+      -- | Annotation for the substructure.
+      bitsSubAnnot :: a
+    } ->
+    RegisterBitsDecl stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents different ways to refer to register bits, either as an array,
@@ -424,85 +554,122 @@ data RegisterBitsDecl stage f a where
 data RegisterBitsTypeRef stage f a where
   -- | An array of bits with a specified size.
   RegisterBitsArray ::
-    { bitsArrayTypeRef :: RegisterBitsTypeRef stage f a, -- ^ Reference to the array type.
-      bitsArraySize :: Expression stage f a, -- ^ Size of the array.
-      bitsArrayAnnot :: a -- ^ Annotation for the array.
-    } -> RegisterBitsTypeRef stage f a
+    { -- | Reference to the array type.
+      bitsArrayTypeRef :: RegisterBitsTypeRef stage f a,
+      -- | Size of the array.
+      bitsArraySize :: Expression stage f a,
+      -- | Annotation for the array.
+      bitsArrayAnnot :: a
+    } ->
+    RegisterBitsTypeRef stage f a
   -- | A reference to another type by name.
   RegisterBitsReference ::
-    { bitsRefName :: Name f a, -- ^ The name of the referenced type.
-      bitsRefAnnot :: a -- ^ Annotation for the reference.
-    } -> RegisterBitsTypeRef stage f a
-  -- | An anonymous type for register bits, used in Stage1.
+    { -- | The name of the referenced type.
+      bitsRefName :: Name f a,
+      -- | Annotation for the reference.
+      bitsRefAnnot :: a
+    } ->
+    RegisterBitsTypeRef stage f a
+  -- | An anonymous type for register bits, used in Parsed.
   RegisterBitsAnonymousType ::
-    { anonBitsWitness :: Witness (stage == Stage1), -- ^ Witness for stage constraint.
-      anonBitsType :: AnonymousBitsType stage f a, -- ^ The anonymous type.
-      anonBitsAnnot :: a -- ^ Annotation for the anonymous type.
-    } -> RegisterBitsTypeRef stage f a
+    { -- | Witness for stage constraint.
+      anonBitsWitness :: Witness (stage == Parsed),
+      -- | The anonymous type.
+      anonBitsType :: AnonymousBitsType stage f a,
+      -- | Annotation for the anonymous type.
+      anonBitsAnnot :: a
+    } ->
+    RegisterBitsTypeRef stage f a
   -- | A direct specification of bits as an expression.
   RegisterBitsJustBits ::
-    { justBitsExpr :: Expression stage f a, -- ^ Expression for the bits.
-      justBitsAnnot :: a -- ^ Annotation for the bits.
-    } -> RegisterBitsTypeRef stage f a
+    { -- | Expression for the bits.
+      justBitsExpr :: Expression stage f a,
+      -- | Annotation for the bits.
+      justBitsAnnot :: a
+    } ->
+    RegisterBitsTypeRef stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- | Represents an anonymous bit type, such as an enum, used in Stage1.
+-- | Represents an anonymous bit type, such as an enum, used in Parsed.
 data AnonymousBitsType stage f a where
   AnonymousEnumBody ::
-    { anonEnumExpr :: Expression stage f a, -- ^ Expression defining the enum size.
-      anonEnumBody :: f (EnumBody stage f a), -- ^ The body of the enum.
-      anonEnumAnnot :: a -- ^ Annotation for the anonymous enum.
-    } -> AnonymousBitsType stage f a
+    { -- | Expression defining the enum size.
+      anonEnumExpr :: Expression stage f a,
+      -- | The body of the enum.
+      anonEnumBody :: f (EnumBody stage f a),
+      -- | Annotation for the anonymous enum.
+      anonEnumAnnot :: a
+    } ->
+    AnonymousBitsType stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents a bit type, either an enumeration or raw bits.
 data BitType (stage :: Stage) (f :: Type -> Type) a where
   -- | An enumeration type for bits.
   EnumBitType ::
-    { enumBitExpr :: Expression stage f a, -- ^ Expression defining the enum size.
-      enumBitBody :: f (EnumBody stage f a), -- ^ The body of the enum.
-      enumBitAnnot :: a -- ^ Annotation for the enumeration.
-    } -> BitType stage f a
+    { -- | Expression defining the enum size.
+      enumBitExpr :: Expression stage f a,
+      -- | The body of the enum.
+      enumBitBody :: f (EnumBody stage f a),
+      -- | Annotation for the enumeration.
+      enumBitAnnot :: a
+    } ->
+    BitType stage f a
   -- | A raw bit type.
   RawBits ::
-    { rawBitsExpr :: Expression stage f a, -- ^ Expression defining the bits.
-      rawBitsAnnot :: a -- ^ Annotation for the raw bits.
-    } -> BitType stage f a
+    { -- | Expression defining the bits.
+      rawBitsExpr :: Expression stage f a,
+      -- | Annotation for the raw bits.
+      rawBitsAnnot :: a
+    } ->
+    BitType stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents the body of an enumeration.
 data EnumBody (stage :: Stage) (f :: Type -> Type) a where
   EnumBody ::
-    { enumConsts :: [Directed EnumConstantDecl stage f a], -- ^ Enum constant declarations.
-      enumBodyAnnot :: a -- ^ Annotation for the enum body.
-    } -> EnumBody stage f a
+    { -- | Enum constant declarations.
+      enumConsts :: [Directed EnumConstantDecl stage f a],
+      -- | Annotation for the enum body.
+      enumBodyAnnot :: a
+    } ->
+    EnumBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents a declaration for an enumeration constant.
 data EnumConstantDecl stage f a where
   -- | A named constant in the enum.
   EnumConstantDecl ::
-    { enumConstIdent :: Identifier f a, -- ^ Identifier for the constant.
-      enumConstExpr :: Expression stage f a, -- ^ Expression defining the constant.
-      enumConstAnnot :: a -- ^ Annotation for the constant.
-    } -> EnumConstantDecl stage f a
+    { -- | Identifier for the constant.
+      enumConstIdent :: Identifier f a,
+      -- | Expression defining the constant.
+      enumConstExpr :: Expression stage f a,
+      -- | Annotation for the constant.
+      enumConstAnnot :: a
+    } ->
+    EnumConstantDecl stage f a
   -- | A reserved value in the enum.
   EnumConstantReserved ::
-    { enumReservedExpr :: Expression stage f a, -- ^ Expression for the reserved value.
-      enumReservedAnnot :: a -- ^ Annotation for the reserved value.
-    } -> EnumConstantDecl stage f a
+    { -- | Expression for the reserved value.
+      enumReservedExpr :: Expression stage f a,
+      -- | Annotation for the reserved value.
+      enumReservedAnnot :: a
+    } ->
+    EnumConstantDecl stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
 -- | Represents the body of a package, containing a list of declarations.
 data PackageBody (stage :: Stage) (f :: Type -> Type) a where
   PackageBody ::
-    { packageBodyDecls :: [Directed FiddleDecl stage f a], -- ^ Declarations in the package.
-      packageBodyAnnot :: a -- ^ Annotation for the package body.
-    } -> PackageBody stage f a
+    { -- | Declarations in the package.
+      packageBodyDecls :: [Directed FiddleDecl stage f a],
+      -- | Annotation for the package body.
+      packageBodyAnnot :: a
+    } ->
+    PackageBody stage f a
   deriving (Generic, Annotated, Typeable)
 
 deriving instance (Alter (ImportType stage)) => Alter (PackageBody stage)
 
-
 squeeze :: (Alter t, Traversable f, Monad f) => t f a -> f (t Identity a)
 squeeze = alter (fmap Identity) return