Rename some of the stages.

Stage1 -> Parsed
Stage2 -> Expanded
Stage3 -> Checked

1 files changed, 523 insertions, 0 deletions

diff --git a/src/Language/Fiddle/Compiler/ConsistencyCheck.hs b/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
new file mode 100644
index 0000000..90f4aa4
--- /dev/null
+++ b/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
@@ -0,0 +1,523 @@
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE IncoherentInstances #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+module Language.Fiddle.Compiler.ConsistencyCheck (checkConsistency) where
+
+import Control.Monad (forM, forM_, unless, when)
+import Control.Monad.Identity (Identity (Identity))
+import Control.Monad.RWS (MonadState (get, put), MonadWriter (tell), gets, modify')
+import Data.Foldable (Foldable (toList), foldlM)
+import Data.Functor.Identity
+import qualified Data.IntMap as IntMap
+import Data.Kind (Type)
+import Data.List (inits, intercalate)
+import Data.List.NonEmpty (NonEmpty (..))
+import qualified Data.List.NonEmpty as NonEmpty
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Maybe (fromMaybe)
+import qualified Data.Set as Set
+import qualified Data.Text as Text
+import Data.Word (Word32)
+import Language.Fiddle.Ast
+import Language.Fiddle.Compiler
+import Language.Fiddle.Internal.Scopes
+import Language.Fiddle.Types (Commented (unCommented), SourceSpan)
+import Text.Printf (printf)
+import Prelude hiding (unzip)
+
+newtype GlobalState = GlobalState
+  { globalScope :: Scope String (Either SizeBits SizeBytes)
+  }
+
+newtype LocalState = LocalState (ScopePath String)
+
+type I = Identity
+
+type Annot = Commented SourceSpan
+
+type SizeBits = Word32
+
+type SizeBytes = Word32
+
+checkConsistency ::
+  FiddleUnit Expanded I Annot ->
+  Compile () (FiddleUnit Checked I Annot)
+checkConsistency =
+  fmap snd
+    . subCompile (GlobalState mempty)
+    . advanceStage (LocalState mempty)
+
+instance CompilationStage Expanded where
+  type StageAfter Expanded = Checked
+  type StageMonad Expanded = Compile GlobalState
+  type StageState Expanded = LocalState
+  type StageFunctor Expanded = Identity
+  type StageAnnotation Expanded = Commented SourceSpan
+
+deriving instance AdvanceStage Expanded FiddleUnit
+
+deriving instance AdvanceStage Expanded Expression
+
+deriving instance AdvanceStage Expanded ObjType
+
+deriving instance AdvanceStage Expanded DeferredRegisterBody
+
+deriving instance AdvanceStage Expanded RegisterBitsDecl
+
+deriving instance AdvanceStage Expanded RegisterBitsTypeRef
+
+deriving instance AdvanceStage Expanded AnonymousBitsType
+
+deriving instance AdvanceStage Expanded BitType
+
+deriving instance AdvanceStage Expanded EnumBody
+
+deriving instance AdvanceStage Expanded EnumConstantDecl
+
+deriving instance AdvanceStage Expanded PackageBody
+
+deriving instance (AdvanceStage Expanded t) => AdvanceStage Expanded (Directed t)
+
+instance AdvanceStage Expanded RegisterBody where
+  advanceStage s body = fst <$> registerBodyToStage3 s body
+
+instance AdvanceStage Expanded ObjTypeBody where
+  advanceStage s body = fst <$> objTypeBodyToStage3 s body 0
+
+instance AdvanceStage Expanded FiddleDecl where
+  modifyState t s = case t of
+    (BitsDecl id typ a) -> do
+      typeSize <- getTypeSize typ
+      insertTypeSize s id typeSize
+      return s
+    (PackageDecl n _ _) -> do
+      let strs = nameToList n
+      let (LocalState scopePath) = s
+
+      return $
+        LocalState $
+          scopePath {currentScope = strs ++ currentScope scopePath}
+    (UsingDecl n _) ->
+      let (LocalState scopePath) = s
+       in return $
+            LocalState $
+              scopePath
+                { usingPaths = nameToList n : usingPaths scopePath
+                }
+    _ -> return s
+
+nameToList :: Name f a -> [String]
+nameToList (Name idents _) = map (\(Identifier (Text.unpack -> s) _) -> s) (toList idents)
+
+objTypeBodyToStage3 ::
+  LocalState ->
+  ObjTypeBody Expanded I Annot ->
+  Word32 ->
+  Compile GlobalState (ObjTypeBody Checked I Annot, Word32)
+objTypeBodyToStage3 st (ObjTypeBody bodyType decls a) startOff = do
+  let isUnion = case bodyType of
+        Union {} -> True
+        _ -> False
+  (cur, returned) <-
+    foldlM
+      ( \(cursor, returned) decl ->
+          case undirected decl of
+            RegisterDecl mMod mIdent expr mBody a -> do
+              (s3RegisterBody, mCalculatedSize) <-
+                fUnzip <$> mapM (registerBodyToStage3 st) mBody
+
+              nExpr <- advanceStage st expr
+
+              let s3 =
+                    mapDirected
+                      ( const $
+                          RegisterDecl
+                            mMod
+                            mIdent
+                            nExpr
+                            s3RegisterBody
+                            a
+                      )
+                      decl
+
+              declaredSizeBits <- fromIntegral <$> exprToSize expr
+
+              when ((declaredSizeBits `mod` 8) /= 0) $
+                tell
+                  [ Diagnostic
+                      Error
+                      "Register size is not a multiple of 8. Please pad register size to align with 8. "
+                      (unCommented a)
+                  ]
+
+              forM_ mCalculatedSize $ \(fromIntegral -> calculatedSize) ->
+                unless (calculatedSize == declaredSizeBits) $
+                  let helpful =
+                        if calculatedSize < declaredSizeBits
+                          then
+                            printf
+                              "\nPerhaps you should add 'reserved(%d)' to the end of your register declaration?"
+                              (declaredSizeBits - calculatedSize)
+                          else ""
+                   in tell
+                        [ Diagnostic
+                            Error
+                            ( printf
+                                "Calculated size %d does not match declared size %d.%s"
+                                calculatedSize
+                                declaredSizeBits
+                                helpful
+                            )
+                            (unCommented a)
+                        ]
+
+              if isUnion
+                then
+                  checkUnion cursor (declaredSizeBits `div` 8) (s3 : returned) a
+                else
+                  return (cursor + declaredSizeBits `div` 8, s3 : returned)
+            TypeSubStructure (Identity subBody) maybeIdent annot -> do
+              (newBody, size) <-
+                objTypeBodyToStage3
+                  st
+                  subBody
+                  ( if isUnion then startOff else cursor
+                  )
+              let s3 = mapDirected (const $ TypeSubStructure (Identity newBody) maybeIdent annot) decl
+
+              checkTypesSubStructure subBody maybeIdent annot
+              if isUnion
+                then
+                  checkUnion cursor size (s3 : returned) a
+                else
+                  return (cursor + size, s3 : returned)
+            ReservedDecl expr annot -> do
+              size' <- fromIntegral <$> exprToSize expr
+              when ((size' `mod` 8) /= 0) $
+                tell
+                  [ Diagnostic
+                      Error
+                      "Can only reserve a multiple of 8 bits in this context."
+                      (unCommented a)
+                  ]
+
+              expr' <- advanceStage st expr
+              let size = size' `div` 8
+              let s3 = mapDirected (const $ ReservedDecl expr' annot) decl
+              if isUnion
+                then
+                  checkUnion cursor size (s3 : returned) a
+                else
+                  return (cursor + size, s3 : returned)
+            AssertPosStatement _ expr a -> do
+              declaredPos <- fromIntegral <$> exprToSize expr
+
+              let expectedPos = if isUnion then startOff else cursor + startOff
+
+              when (expectedPos /= declaredPos) $ do
+                tell
+                  [ Diagnostic
+                      Error
+                      ( printf
+                          "Position assertion failed. Asserted 0x%x, calculated 0x%x"
+                          declaredPos
+                          expectedPos
+                      )
+                      (unCommented a)
+                  ]
+              return (cursor, returned)
+      )
+      (0, [])
+      decls
+
+  return (ObjTypeBody bodyType (reverse returned) a, cur)
+  where
+    checkTypesSubStructure
+      (ObjTypeBody bodyType decls _)
+      maybeIdent
+      annot =
+        let emitWarning s = tell [Diagnostic Warning s (unCommented annot)]
+         in case () of
+              ()
+                | [_] <- decls,
+                  (Union {}) <- bodyType ->
+                    emitWarning "Union with a single field. Should this be a struct?"
+              ()
+                | [_] <- decls,
+                  (Struct {}) <- bodyType,
+                  Nothing <- maybeIdent ->
+                    emitWarning "Anonymous sub-struct with single field is superfluous."
+              ()
+                | [] <- decls ->
+                    emitWarning
+                      ( printf
+                          "Empty sub-%s is superfluous."
+                          ( case bodyType of
+                              Union {} -> "union"
+                              Struct {} -> "struct"
+                          )
+                      )
+              _ -> return ()
+    fUnzip xs = (fst <$> xs, snd <$> xs)
+    pushApply :: Maybe (a, b) -> (Maybe a, Maybe b)
+    pushApply (Just (a, b)) = (Just a, Just b)
+    pushApply Nothing = (Nothing, Nothing)
+
+registerBodyToStage3 ::
+  LocalState ->
+  RegisterBody Expanded I Annot ->
+  Compile GlobalState (RegisterBody Checked I Annot, Word32)
+registerBodyToStage3
+  st
+  (RegisterBody bodyType (Identity deferredRegisterBody) a') = do
+    let isUnion = case bodyType of
+          Union {} -> True
+          _ -> False
+
+    case deferredRegisterBody of
+      DeferredRegisterBody decls a -> do
+        (cur, returned) <-
+          foldlM
+            ( \(cursor, returned) decl ->
+                case undirected decl of
+                  ReservedBits expr a -> do
+                    size <- fromIntegral <$> exprToSize expr
+                    expr' <- advanceStage st expr
+                    let s3 =
+                          mapDirected
+                            (const $ ReservedBits expr' a)
+                            decl
+                    if isUnion
+                      then checkUnion cursor size (s3 : returned) a
+                      else
+                        return (cursor + size, s3 : returned)
+                  BitsSubStructure registerBody maybeIdent annot -> do
+                    checkBitsSubStructure registerBody maybeIdent annot
+
+                    (newBody, subsize) <- registerBodyToStage3 st registerBody
+                    let s3 =
+                          mapDirected
+                            (const $ BitsSubStructure newBody maybeIdent annot)
+                            decl
+
+                    if isUnion
+                      then checkUnion cursor subsize (s3 : returned) a
+                      else
+                        return (cursor + subsize, s3 : returned)
+                  DefinedBits modifier identifier typeref a -> do
+                    (s3TypeRef, size) <- registerBitsTypeRefToStage3 st typeref
+                    let s3 =
+                          mapDirected
+                            (const $ DefinedBits modifier identifier s3TypeRef a)
+                            decl
+
+                    if isUnion
+                      then checkUnion cursor size (s3 : returned) a
+                      else
+                        return (cursor + size, s3 : returned)
+            )
+            (0, [])
+            decls
+
+        return
+          ( RegisterBody
+              bodyType
+              (Identity (DeferredRegisterBody (reverse returned) a))
+              a',
+            cur
+          )
+    where
+      checkBitsSubStructure
+        (RegisterBody bodyType (Identity (DeferredRegisterBody decls _)) _)
+        maybeIdent
+        annot =
+          let emitWarning s = tell [Diagnostic Warning s (unCommented annot)]
+           in case () of
+                ()
+                  | [_] <- decls,
+                    (Union {}) <- bodyType ->
+                      emitWarning "Union with a single field. Should this be a struct?"
+                ()
+                  | [_] <- decls,
+                    (Struct {}) <- bodyType,
+                    Nothing <- maybeIdent ->
+                      emitWarning "Anonymous sub-struct with single field is superfluous."
+                ()
+                  | [] <- decls ->
+                      emitWarning
+                        ( printf
+                            "Empty sub-%s is superfluous."
+                            ( case bodyType of
+                                Union {} -> "union"
+                                Struct {} -> "struct"
+                            )
+                        )
+                _ -> return ()
+
+registerBitsTypeRefToStage3 ::
+  LocalState ->
+  RegisterBitsTypeRef Expanded I Annot ->
+  Compile GlobalState (RegisterBitsTypeRef Checked I Annot, Word32)
+registerBitsTypeRefToStage3 localState = \case
+  RegisterBitsArray ref expr a -> do
+    (ref', size) <- registerBitsTypeRefToStage3 localState ref
+    multiplier <- exprToSize expr
+    expr' <- advanceStage localState expr
+    return
+      ( RegisterBitsArray ref' expr' a,
+        size * fromIntegral multiplier
+      )
+  RegisterBitsReference name a ->
+    (RegisterBitsReference name a,) <$> lookupTypeSize localState name
+  RegisterBitsJustBits expr a -> do
+    expr' <- advanceStage localState expr
+    (RegisterBitsJustBits expr' a,)
+      . fromIntegral
+      <$> exprToSize expr
+
+checkUnion :: Word32 -> Word32 -> b -> Commented SourceSpan -> Compile s (Word32, b)
+checkUnion cursor subsize ret a = do
+  when (cursor /= 0 && subsize /= cursor) $ do
+    tell
+      [ Diagnostic
+          Warning
+          ( printf
+              "Jagged union found. Found size %d, expected %d.\n \
+              \ Please wrap smaller fields in a struct with padding so all \
+              \ fields are the same size?"
+              subsize
+              cursor
+          )
+          (unCommented a)
+      ]
+  return (max cursor subsize, ret)
+
+exprToSize ::
+  (NumberType stage ~ Integer) =>
+  Expression stage I Annot ->
+  Compile s Integer
+exprToSize (LitNum num _) = return num
+exprToSize e = do
+  tell [Diagnostic Error "Variables not allowed" (unCommented $ annot e)]
+  compilationFailure
+
+lookupTypeSize :: LocalState -> Name I Annot -> Compile GlobalState SizeBits
+lookupTypeSize (LocalState scopePath) (Name idents a) = do
+  -- Convert the list of identifiers to a string path
+  let path = fmap (\(Identifier s _) -> Text.unpack s) idents
+
+  -- Get the current scope and perform the lookup
+  results <- gets $ lookupScopeWithPath scopePath path . globalScope
+
+  case results of
+    -- Successfully resolved to a unique size
+    [(_, Right sz)] -> return sz
+    -- Multiple ambiguous results found
+    matches@(_ : _) -> do
+      -- Generate a list of ambiguous paths for error reporting
+      let ambiguousPaths =
+            map
+              ( \(resolvedPath, _) ->
+                  intercalate "." (NonEmpty.toList resolvedPath)
+              )
+              matches
+      tell
+        [ Diagnostic
+            Error
+            ( printf
+                "Ambiguous occurrence of '%s'. Multiple matches found:\n%s"
+                (intercalate "." $ NonEmpty.toList path)
+                (unlines ambiguousPaths) -- List all ambiguous paths
+            )
+            (unCommented a)
+        ]
+      compilationFailure
+
+    -- No matches found
+    _ -> do
+      tell
+        [ Diagnostic
+            Error
+            ( printf
+                "Cannot resolve '%s'. No matching symbols found."
+                (intercalate "." $ NonEmpty.toList path)
+            )
+            (unCommented a)
+        ]
+      compilationFailure
+
+getTypeSize :: BitType Expanded I Annot -> Compile s SizeBits
+getTypeSize (RawBits expr _) = fromIntegral <$> exprToSize expr
+getTypeSize (EnumBitType expr (Identity (EnumBody constants _)) ann) = do
+  declaredSize <- fromIntegral <$> exprToSize expr
+
+  -- If the declared size is less than or equal to 4, we'll enforce that the
+  -- enum is packed. This is to make sure the user has covered all bases.
+  when (declaredSize <= 4) $ do
+    imap <-
+      foldlM
+        ( \imap (undirected -> enumConst) -> do
+            number <- case enumConst of
+              EnumConstantDecl _ expr _ -> exprToSize expr
+              EnumConstantReserved expr _ -> exprToSize expr
+
+            when (number >= 2 ^ declaredSize) $
+              tell
+                [ Diagnostic
+                    Error
+                    ( printf
+                        "Enum constant too large. Max allowed %d\n"
+                        ((2 :: Int) ^ declaredSize)
+                    )
+                    (unCommented (annot enumConst))
+                ]
+
+            return $ IntMap.insert (fromIntegral number) True imap
+        )
+        IntMap.empty
+        constants
+    let missing =
+          filter (not . (`IntMap.member` imap)) [0 .. 2 ^ declaredSize - 1]
+    unless (null missing) $
+      tell
+        [ Diagnostic
+            Warning
+            ( printf
+                "Missing enum constants %s. Small enums should be fully \
+                \ populated. Use 'reserved' if needed."
+                (intercalate ", " (map show missing))
+            )
+            (unCommented ann)
+        ]
+
+  return declaredSize
+
+diagnosticError :: String -> Annot -> Compile a ()
+diagnosticError str a = tell [Diagnostic Error str (unCommented a)]
+
+insertTypeSize ::
+  LocalState ->
+  Identifier f Annot ->
+  SizeBits ->
+  Compile GlobalState ()
+insertTypeSize (LocalState scopePath) (Identifier s annot) size = do
+  modifyM $
+    \(GlobalState globalScope) ->
+      let fullName =
+            NonEmpty.prependList
+              (currentScope scopePath)
+              (NonEmpty.singleton (Text.unpack s))
+       in case upsertScope fullName (Right size) globalScope of
+            (Just _, _) -> do
+              diagnosticError (printf "Duplicate type %s" s) annot
+              compilationFailure
+            (Nothing, n) -> return $ GlobalState n
+  where
+    modifyM fn = do
+      s <- get
+      put =<< fn s