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

-- | Qualification compilation phase.
--
-- The qualification phase is responsible for resolving all type references in
-- the AST to their fully-qualified counterparts. This process involves
-- replacing unqualified references with their fully-qualified names and
-- attaching the necessary metadata to each reference. This enriched information
-- is then available for use in later stages of the compilation pipeline.
--
-- In this phase, symbol resolution statements (such as 'using' statements) are
-- removed, as they become unnecessary once references are fully qualified.
module Language.Fiddle.Compiler.Qualification (qualificationPhase) where

import Control.Monad.RWS (MonadWriter (tell))
import Control.Monad.State
import Data.Foldable (foldlM, toList)
import Data.List (intercalate)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (isNothing, mapMaybe)
import qualified Data.Text
import Data.Word
import Language.Fiddle.Ast
import Language.Fiddle.Compiler
import Language.Fiddle.Compiler.ConsistencyCheck ()
import Language.Fiddle.Internal.Scopes
import Language.Fiddle.Internal.UnitInterface as UnitInterface
import Language.Fiddle.Internal.UnitNumbers
import Language.Fiddle.Types
import Text.Printf (printf)

type S = Expanded

data GlobalState = GlobalState
  { unitInterface :: UnitInterface,
    uniqueCounter :: Int
  }

data LocalState = LocalState
  { currentScopePath :: ScopePath String,
    ephemeralScope :: Scope String (Metadata, ExportedDecl)
  }

type F = Either [Diagnostic]

type A = Commented SourceSpan

type M = Compile GlobalState

uniqueString :: String -> M String
uniqueString prefix = do
  cnt <- gets uniqueCounter
  modify $ \g -> g {uniqueCounter = cnt + 1}
  return $ "_" ++ prefix ++ show cnt

uniqueIdentifier :: String -> a -> M (Identifier F a)
uniqueIdentifier prefix a = (\s -> Identifier (Data.Text.pack s) a) <$> uniqueString prefix

instance CompilationStage Expanded where
  type StageAfter Expanded = Qualified
  type StageMonad Expanded = M
  type StageState Expanded = LocalState
  type StageFunctor Expanded = F
  type StageAnnotation Expanded = A

qualificationPhase :: CompilationPhase Expanded Qualified
qualificationPhase =
  pureCompilationPhase $ \t -> do
    raw <-
      fmap snd $
        subCompile (GlobalState mempty 0) $
          advanceStage
            (LocalState mempty mempty)
            (soakA t)

    squeezeDiagnostics raw

pushIdent :: Identifier f a -> LocalState -> LocalState
pushIdent i = pushIdents [i]

pushIdents :: (Foldable t) => t (Identifier f a) -> LocalState -> LocalState
pushIdents =
  ( \case
      [] -> id
      (i : is) ->
        modifyCurrentScopePath (pushScope $ fmap identToString (i :| is))
  )
    . toList

instance
  StageConvertible
    Expanded
    (When False String)
    (When True String)
  where
  convertInStage _ _ _ _ = Present <$> uniqueString "reserved"

deriving instance AdvanceStage S ObjTypeBody

deriving instance AdvanceStage S DeferredRegisterBody

deriving instance AdvanceStage S RegisterBody

deriving instance AdvanceStage S AnonymousBitsType

deriving instance AdvanceStage S ImportStatement

deriving instance AdvanceStage S BitType

deriving instance AdvanceStage S EnumBody

deriving instance AdvanceStage S EnumConstantDecl

instance AdvanceStage S RegisterBitsDecl where
  advanceStage localState = \case
    ReservedBits expr an -> ReservedBits <$> advanceStage localState expr <*> pure an
    BitsSubStructure bod name an ->
      BitsSubStructure
        <$> advanceStage localState bod
        <*> pure name
        <*> pure an
    DefinedBits _ mod ident typ an -> do
      let qMeta =
            QBitsMetadata
              { bitsSpan = Vacant,
                bitsFullPath =
                  qualifyPath
                    (currentScopePath localState)
                    (NonEmpty.singleton (identToString ident))
              }
      DefinedBits (Present qMeta) mod ident
        <$> advanceStage localState typ
        <*> pure an

instance AdvanceStage S ObjTypeDecl where
  advanceStage localState = \case
    AssertPosStatement d e a ->
      AssertPosStatement d <$> advanceStage localState e <*> pure a
    RegisterDecl _ mod ident size bod ann -> do
      ident' <- guaranteeM (uniqueIdentifier "reg" ann) ident
      let localState' = pushIdents ident localState

      let qRegMeta =
            QRegMetadata
              { regSpan = Vacant,
                regIsPadding = False,
                regIsUnnamed = isNothing (toMaybe ident),
                regFullPath =
                  qualifyPath
                    (currentScopePath localState)
                    (NonEmpty.singleton (identToString (unwrap ident')))
              }

      RegisterDecl
        (Present qRegMeta)
        (guarantee (ModifierKeyword Rw ann) mod)
        ident'
        <$> advanceStage localState' size
        <*> mapM (advanceStage localState') bod
        <*> pure ann
    ReservedDecl _ expr ann -> do
      ident <- uniqueIdentifier "reserved" ann

      let qRegMeta =
            QRegMetadata
              { regSpan = Vacant,
                regIsPadding = True,
                regIsUnnamed = True,
                regFullPath =
                  qualifyPath
                    (currentScopePath localState)
                    (NonEmpty.singleton (identToString ident))
              }

      RegisterDecl
        (Present qRegMeta)
        (Guaranteed $ ModifierKeyword Pr ann)
        (Guaranteed ident)
        <$> advanceStage localState expr
        <*> pure Nothing
        <*> pure ann
    TypeSubStructure bod name an -> do
      let localState' = pushIdents name localState
      TypeSubStructure
        <$> mapM (advanceStage localState') bod
        <*> pure name
        <*> pure an

deriving instance AdvanceStage S (Expression u)

instance AdvanceStage S RegisterBitsTypeRef where
  advanceStage localState = \case
    RegisterBitsArray a b c ->
      RegisterBitsArray
        <$> advanceStage localState a
        <*> advanceStage localState b
        <*> pure c
    RegisterBitsJustBits a b ->
      RegisterBitsJustBits
        <$> advanceStage localState a
        <*> pure b
    RegisterBitsReference _ name a -> do
      v <- fmap (Present . snd) <$> resolveName name localState
      return $ RegisterBitsReference v name a

instance AdvanceStage S ObjType where
  advanceStage localState = \case
    ArrayObjType a b c ->
      ArrayObjType
        <$> advanceStage localState a
        <*> advanceStage localState b
        <*> pure c
    ReferencedObjType _ name a -> do
      v <- fmap (Present . snd) <$> resolveName name localState
      return $ ReferencedObjType v name a

deriving instance (AdvanceStage S t) => AdvanceStage S (Directed t)

instance AdvanceStage S PackageBody where
  advanceStage localState (PackageBody decls a) =
    PackageBody <$> advanceFiddleDecls localState decls <*> pure a

instance AdvanceStage S FiddleUnit where
  advanceStage localState (FiddleUnit v decls a) =
    FiddleUnit v <$> advanceFiddleDecls localState decls <*> pure a

modifyEphemeralScope ::
  ( Scope String (Metadata, ExportedDecl) -> Scope String (Metadata, ExportedDecl)
  ) ->
  LocalState ->
  LocalState
modifyEphemeralScope fn ls@LocalState {ephemeralScope = es} =
  ls {ephemeralScope = fn es}

modifyCurrentScopePath ::
  (ScopePath String -> ScopePath String) ->
  LocalState ->
  LocalState
modifyCurrentScopePath fn ls@LocalState {currentScopePath = cs} =
  ls {currentScopePath = fn cs}

resolveIdent :: (ExportableDecl d, Functor f) => Identifier f A -> LocalState -> M (F ([String], d))
resolveIdent i = resolveSymbol (annot i) [identToString i]

resolveName :: (ExportableDecl d, Functor f) => Name f A -> LocalState -> M (F ([String], d))
resolveName n = resolveSymbol (annot n) (toList $ nameToList n)

resolveSymbol :: (ExportableDecl d) => A -> [String] -> LocalState -> M (F ([String], d))
resolveSymbol a (p : ps) (LocalState {ephemeralScope = ephemeralScope, currentScopePath = currentPath}) = do
  GlobalState {unitInterface = UnitInterface {rootScope = rootScope}} <- get

  let matches =
        concatMap
          ( mapMaybe (\(p, (m, e)) -> (p,) . (m,) <$> fromExportedDecl e)
              . lookupScopeWithPath currentPath (p :| ps)
          )
          [rootScope, ephemeralScope]

  return $
    case matches of
      [(p, (_, e))] -> Right (toList p, e)
      [] ->
        Left
          [ Diagnostic
              Error
              ( printf "Could not resolve symbol %s" (intercalate "." (p : ps))
              )
              (unCommented a)
          ]
      (_ : _ : _) -> do
        Left
          [ Diagnostic
              Error
              ( printf
                  "Ambiguous occurance of %s"
                  (intercalate "." (p : ps))
              )
              (unCommented a)
          ]
resolveSymbol a _ _ =
  return $ Left [Diagnostic Error "Empty path provided (this is a bug)" (unCommented a)]

qMd :: (Applicative f) => t -> f (QMd Qualified t)
qMd = pure . Present

advanceFiddleDecls ::
  LocalState ->
  [Directed FiddleDecl S F A] ->
  M [Directed FiddleDecl Qualified F A]
advanceFiddleDecls localState decls = fmap (reverse . fst) $ do
  foldlM
    ( \(declsRet, localState' :: LocalState) unsqeezedd -> do
        d <- case squeeze unsqeezedd of
          Left diags -> tell diags >> compilationFailure
          Right x -> return x
        case unsqeezedd of
          (Directed directives t dann) ->
            let doReturn ::
                  FiddleDecl Qualified F A ->
                  M ([Directed FiddleDecl Qualified F A], LocalState)
                doReturn v = return (Directed directives v dann : declsRet, localState')
                doReturnWith s v = return (Directed directives v dann : declsRet, s)
                qualify = qualifyPath (currentScopePath localState')
                metadata = directiveToMetadata d
             in case t of
                  UsingDecl {usingName = name} ->
                    return (declsRet, modifyCurrentScopePath (addUsingPath (nameToList name)) localState')
                  OptionDecl key value ann -> doReturn $ OptionDecl key value ann
                  ImportDecl st@(ImportStatement {importInterface = interface}) a ->
                    let localState'' = modifyEphemeralScope (<> rootScope (unwrap interface)) localState'
                     in doReturnWith localState''
                          =<< ImportDecl
                            <$> advanceStage localState'' st
                            <*> pure a
                  PackageDecl _ name body ann ->
                    let qualifiedName = qualify (nameToList name)
                        localState'' = modifyCurrentScopePath (pushScope (nameToList name)) localState'
                        decl = ExportedPackageDecl (metadata qualifiedName)
                     in do
                          insertDecl decl
                          doReturn
                            =<< PackageDecl
                              (qMd decl)
                              name
                              <$> mapM (advanceStage localState'') body
                              <*> pure ann
                  LocationDecl _ ident expr ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                     in do
                          exprValue <- expressionToIntM expr
                          let decl =
                                ExportedLocationDecl
                                  (metadata qualifiedName)
                                  exprValue
                          insertDecl decl
                          doReturn
                            =<< LocationDecl
                              (qMd decl)
                              ident
                              <$> advanceStage localState' expr
                              <*> pure ann
                  BitsDecl _ ident typ ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                     in do
                          sizeBits <- getBitTypeDeclaredSize typ
                          let decl =
                                ExportedBitsDecl
                                  (metadata qualifiedName)
                                  sizeBits
                          insertDecl decl
                          doReturn
                            =<< BitsDecl
                              (qMd decl)
                              ident
                              <$> advanceStage localState' typ
                              <*> pure ann
                  ObjTypeDecl _ ident body ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                        localState'' = modifyCurrentScopePath (pushScope (NonEmpty.singleton $ identToString ident)) localState'
                     in do
                          typeSize <- calculateTypeSize =<< resolveOrFail body
                          let decl =
                                ExportedTypeDecl
                                  (metadata qualifiedName)
                                  typeSize
                          insertDecl decl
                          doReturn
                            =<< ObjTypeDecl
                              (qMd decl)
                              ident
                              <$> mapM (advanceStage localState'') body
                              <*> pure ann
                  ObjectDecl _ ident loc typ ann ->
                    let qualifiedName = qualify (NonEmpty.singleton (identToString ident))
                     in do
                          location <- resolveLocationExpression localState' loc
                          exportedType <- objTypeToExport localState' typ
                          let decl =
                                ExportedObjectDecl
                                  (metadata qualifiedName)
                                  location
                                  exportedType
                          insertDecl decl
                          doReturn
                            =<< ObjectDecl
                              (qMd decl)
                              ident
                              <$> advanceStage localState' loc
                              <*> advanceStage localState' typ
                              <*> pure ann
    )
    ([], localState)
    decls

insertDecl :: (ExportableDecl d) => d -> M ()
insertDecl decl =
  modify $ \(GlobalState ui c) -> GlobalState (UnitInterface.insert decl ui) c

objTypeToExport :: LocalState -> ObjType Expanded F A -> M ReferencedObjectType
objTypeToExport ls = \case
  ArrayObjType {arraySize = size, arrayObjType = objType} ->
    ArrayObjectType
      <$> objTypeToExport ls objType
      <*> expressionToIntM size
  ReferencedObjType {refName = n} -> do
    (full, _ :: ExportedTypeDecl) <- resolveOrFail =<< resolveName n ls
    case full of
      (f:fs) -> return $ ReferencedObjectType (f :| fs)
      _ -> compilationFailure

calculateTypeSize :: ObjTypeBody Expanded F A -> M (N Bytes)
calculateTypeSize (ObjTypeBody bodyType decls _) =
  ( case bodyType of
      Union {} -> maximum
      Struct {} -> sum
  )
    <$> mapM calculateDeclSize decls
  where
    calculateDeclSize :: Directed ObjTypeDecl Expanded F A -> M (N Bytes)
    calculateDeclSize (undirected -> decl) =
      case decl of
        AssertPosStatement {} -> return 0
        RegisterDecl {regSize = size} -> fst . bitsToBytes <$> expressionToIntM size
        ReservedDecl {reservedExpr = size} -> fst . bitsToBytes <$> expressionToIntM size
        TypeSubStructure {subStructureBody = b} -> calculateTypeSize =<< resolveOrFail b

getBitTypeDeclaredSize :: BitType Expanded F A -> M (N Bits)
getBitTypeDeclaredSize = \case
  (EnumBitType declaredSize _ _) -> expressionToIntM declaredSize
  (RawBits declaredSize _) -> expressionToIntM declaredSize

resolveLocationExpression ::
  (stage .< Expanded ~ False) =>
  LocalState ->
  Expression u stage F A ->
  M (N u)
resolveLocationExpression ls (Var var _) = do
  (_, ExportedLocationDecl _ v) <- resolveOrFail =<< resolveName var ls
  return (fromIntegral v)
resolveLocationExpression _ e = expressionToIntM e

expressionToIntM ::
  (stage .< Expanded ~ False) =>
  Expression u stage f A ->
  M (N u)
expressionToIntM expr =
  resolveOrFail $
    either
      ( \reason -> Left [Diagnostic Error reason (unCommented $ annot expr)]
      )
      return
      (expressionToInt expr)