path: root/src/Language/Fiddle/Compiler/ConsistencyCheck.hs

{-# 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