1 files changed, 6 insertions, 4 deletions
diff --git a/src/Language/Fiddle/Compiler/ConsistencyCheck.hs b/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
index e081122..3286e86 100644
--- a/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
+++ b/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
@@ -155,8 +155,10 @@ instance AdvanceStage Expanded RegisterBody where
 instance AdvanceStage Expanded ObjTypeBody where
   advanceStage s body = fst <$> objTypeBodyToStage3 s body 0
 
-instance AdvanceStage Expanded FiddleDecl where
-  modifyState t s = case t of
+deriving instance AdvanceStage Expanded FiddleDecl
+
+instance AdvanceStage Expanded (Directed FiddleDecl) where
+  modifyState (Directed directives t _) s = case t of
     (BitsDecl id typ annotation) -> do
       typeSize <- getTypeSize typ
       insertTypeSize annotation s id typeSize
@@ -177,7 +179,7 @@ instance AdvanceStage Expanded FiddleDecl where
                 }
     _ -> return s
 
-  customAdvanceStage t (LocalState scopePath) = case t of
+  customAdvanceStage (Directed directives t a) (LocalState scopePath) = case t of
     (ObjTypeDecl ident (Identity body) annot) -> do
       (body', size) <- objTypeBodyToStage3 (LocalState scopePath) body 0
 
@@ -190,7 +192,7 @@ instance AdvanceStage Expanded FiddleDecl where
       let ui' = insertIntoUnitInterface fullName ui annot (ExportedObjType size)
       modify' $ \gs -> gs {unitInterface = ui'}
 
-      return $ Just $ ObjTypeDecl ident (Identity body') annot
+      return $ Just $ Directed directives (ObjTypeDecl ident (Identity body') annot) a
     _ -> return Nothing
 
 nameToList :: Name f a -> [String]