1 files changed, 45 insertions, 29 deletions

diff --git a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
index c37be87..706a178 100644
--- a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
+++ b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
@@ -51,6 +51,7 @@ module Language.Fiddle.Ast.Internal.SyntaxTree
   )
 where
 
+import Control.Monad (forM_)
 import Data.Coerce
 import Data.Functor.Identity
 import Data.Kind (Type)
@@ -70,6 +71,9 @@ import Language.Fiddle.Ast.Internal.Kinds
 import Language.Fiddle.Ast.Internal.Stage
 import Language.Fiddle.Internal.UnitInterface (UnitInterface)
 
+type family FiddleUnitInterface (s :: Stage) :: Type where
+  FiddleUnitInterface s = If (s < Checked) () UnitInterface
+
 -- | The Type of number during each stage of compilation. When in the first stage,
 -- numbers are just strings like anything else. In later stages, numbers get
 -- parsed into actual integers. This makes it easier to process later.
@@ -103,7 +107,7 @@ type family WitnessType (s :: Bool) where
 -- to different packages.
 data Name :: SynTree where
   Name :: NonEmpty (Identifier f a) -> a -> Name f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents a directive in the Fiddle language. A directive provides
 --   additional metadata or instructions that the compiler can use during
@@ -117,7 +121,7 @@ data Directive :: SynTree where
       directiveAnnot :: a
     } ->
     Directive f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents the body of a directive, which consists of multiple elements.
 data DirectiveBody :: SynTree where
@@ -128,7 +132,7 @@ data DirectiveBody :: SynTree where
       directiveBodyAnnot :: a
     } ->
     DirectiveBody f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents an element in a directive. Can be either a key or a key-value
 --   pair.
@@ -157,7 +161,7 @@ data DirectiveElement :: SynTree where
       directiveKeyValueAnnot :: a
     } ->
     DirectiveElement f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents expressions that can be used within a directive, either a
 --   string or a number.
@@ -176,7 +180,7 @@ data DirectiveExpression f a where
       directiveNumberAnnot :: a
     } ->
     DirectiveExpression f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | A type that wraps another syntax tree and applies a list of directives to
 --   it.
@@ -192,6 +196,18 @@ data Directed t stage f a where
     Directed t stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+instance
+  (Typeable (Directed t stage), Walk (t stage)) =>
+  Walk (Directed t stage)
+  where
+  walk fn (Directed directives subtree _) s = do
+    s' <- fn subtree s
+    walk fn subtree s'
+
+    forM_ directives $ \d -> do
+      s' <- fn d s
+      walk fn d s'
+
 -- | Apply a function to the underlying subtree in a 'Directed' type.
 mapDirected :: (t s f a -> t' s' f a) -> Directed t s f a -> Directed t' s' f a
 mapDirected fn (Directed dr tfa a) = Directed dr (fn tfa) a
@@ -219,11 +235,14 @@ data FiddleUnit (stage :: Stage) (f :: Type -> Type) a where
   FiddleUnit ::
     { -- | List of declarations.
       fiddleDecls :: [Directed FiddleDecl stage f a],
+      -- | The interface for this FiddleUnit. Early on, this is just () because
+      -- not enough information is provided to determine the interface..
+      fiddleUnitInterface :: FiddleUnitInterface stage,
       -- | Annotation for the 'FiddleUnit'.
       fiddleUnitAnnot :: a
     } ->
     FiddleUnit stage f a
-  deriving (Generic, Annotated, Typeable, Alter)
+  deriving (Generic, Annotated, Typeable, Alter, Walk)
 
 -- | Represents an identifier with an associated annotation.
 data Identifier f a = Identifier
@@ -232,7 +251,7 @@ data Identifier f a = Identifier
     -- | Annotation for the identifier.
     identifierAnnot :: a
   }
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Expressions used within Fiddle, including literals and variables.
 data Expression (s :: Stage) :: SynTree where
@@ -252,7 +271,7 @@ data Expression (s :: Stage) :: SynTree where
       varAnnot :: a
     } ->
     Expression stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents an import statement in the Fiddle language.
 data ImportStatement stage f a where
@@ -261,14 +280,12 @@ data ImportStatement stage f a where
       importPath :: Text,
       -- | Optional list of imported items.
       importList :: Maybe (ImportList f a),
-
       importInterface :: ImportInterface stage,
-
       -- | Annotation for the import statement.
       importStatementAnnot :: a
     } ->
     ImportStatement stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | A list of imported identifiers.
 data ImportList f a where
@@ -279,7 +296,7 @@ data ImportList f a where
       importListAnnot :: a
     } ->
     ImportList f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents top-level declarations in Fiddle.
 data FiddleDecl :: StagedSynTree where
@@ -297,12 +314,11 @@ data FiddleDecl :: StagedSynTree where
   ImportDecl ::
     { -- | The imported type.
       importStatement :: ImportStatement stage f a,
+      -- \| Annotation for the import declaration.
 
       -- | The interface for this imported file. This type depends on the stage
       -- of compilation. Initially it's just '()', but will eventually be resolved
       -- into a 'UnitInterface'.
-
-      -- | Annotation for the import declaration.
       importDeclAnnot :: a
     } ->
     FiddleDecl stage f a
@@ -366,7 +382,7 @@ data FiddleDecl :: StagedSynTree where
       objectAnnot :: a
     } ->
     FiddleDecl stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents the body of an object type, containing a body type (struct or
 --   union), a list of object declarations, and an annotation.
@@ -380,7 +396,7 @@ data ObjTypeBody (stage :: Stage) (f :: Type -> Type) a where
       objBodyAnnot :: a
     } ->
     ObjTypeBody stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents an object type, which can be anonymous, an array, or a
 --   reference to another type.
@@ -413,7 +429,7 @@ data ObjType stage f a where
       refAnnot :: a
     } ->
     ObjType stage f a
-  deriving (Typeable, Generic, Alter, Annotated, Typeable)
+  deriving (Typeable, Generic, Alter, Annotated, Typeable, Walk)
 
 -- | Represents a declaration inside an object type, such as a register, an
 --   assertion, or a substructure.
@@ -460,7 +476,7 @@ data ObjTypeDecl stage f a where
       subStructureAnnot :: a
     } ->
     ObjTypeDecl stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents a modifier for registers (e.g., read-only, read-write).
 data Modifier f a where
@@ -471,7 +487,7 @@ data Modifier f a where
       modifierAnnot :: a
     } ->
     Modifier f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Enumerates the different types of register modifiers.
 data ModifierKeyword = Rw | Ro | Wo
@@ -487,7 +503,7 @@ data DeferredRegisterBody stage f a where
       deferredAnnot :: a
     } ->
     DeferredRegisterBody stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents the body type (struct or union) in an object.
 data BodyType (f :: Type -> Type) a where
@@ -501,7 +517,7 @@ data BodyType (f :: Type -> Type) a where
       structAnnot :: a
     } ->
     BodyType f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents a register body with a body type and deferred bit declarations.
 data RegisterBody stage f a where
@@ -514,7 +530,7 @@ data RegisterBody stage f a where
       regBodyAnnot :: a
     } ->
     RegisterBody stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents declarations within a register, such as defined bits,
 --   reserved bits, or substructures.
@@ -549,7 +565,7 @@ data RegisterBitsDecl stage f a where
       bitsSubAnnot :: a
     } ->
     RegisterBitsDecl stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents different ways to refer to register bits, either as an array,
 --   a reference to a type, an anonymous type, or just bits.
@@ -590,7 +606,7 @@ data RegisterBitsTypeRef stage f a where
       justBitsAnnot :: a
     } ->
     RegisterBitsTypeRef stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents an anonymous bit type, such as an enum, used in Parsed.
 data AnonymousBitsType stage f a where
@@ -603,7 +619,7 @@ data AnonymousBitsType stage f a where
       anonEnumAnnot :: a
     } ->
     AnonymousBitsType stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents a bit type, either an enumeration or raw bits.
 data BitType (stage :: Stage) (f :: Type -> Type) a where
@@ -625,7 +641,7 @@ data BitType (stage :: Stage) (f :: Type -> Type) a where
       rawBitsAnnot :: a
     } ->
     BitType stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents the body of an enumeration.
 data EnumBody (stage :: Stage) (f :: Type -> Type) a where
@@ -636,7 +652,7 @@ data EnumBody (stage :: Stage) (f :: Type -> Type) a where
       enumBodyAnnot :: a
     } ->
     EnumBody stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents a declaration for an enumeration constant.
 data EnumConstantDecl stage f a where
@@ -658,7 +674,7 @@ data EnumConstantDecl stage f a where
       enumReservedAnnot :: a
     } ->
     EnumConstantDecl stage f a
-  deriving (Generic, Annotated, Alter, Typeable)
+  deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents the body of a package, containing a list of declarations.
 data PackageBody (stage :: Stage) (f :: Type -> Type) a where
@@ -669,7 +685,7 @@ data PackageBody (stage :: Stage) (f :: Type -> Type) a where
       packageBodyAnnot :: a
     } ->
     PackageBody stage f a
-  deriving (Generic, Annotated, Typeable, Alter)
+  deriving (Generic, Annotated, Typeable, Alter, Walk)
 
 squeeze :: (Alter t, Traversable f, Monad f) => t f a -> f (t Identity a)
 squeeze = alter (fmap Identity) return