Types.hs

-- SPDX-FileCopyrightText: 2025-2026 Sören Tempel <soeren+git@soeren-tempel.net>
--
-- SPDX-License-Identifier: GPL-3.0-only

module Language.QBE.Types
  ( -- * Identifiers
    UserIdent (..),
    LocalIdent (..),
    BlockIdent (..),
    GlobalIdent (..),

    -- * Types
    BaseType (..),
    baseTypeByteSize,
    baseTypeBitSize,
    ExtType (..),
    extTypeBitSize,
    extTypeByteSize,
    SubWordType (..),
    SubType (..),
    LoadType (..),
    loadByteSize,

    -- * Values
    Const (..),
    DynConst (..),
    Value (..),

    -- * Definitions
    TypeDef (..),
    DataDef (..),
    Linkage (..),
    Field,
    AggType (..),
    dataSize,
    DataObj (..),
    objAlign,
    objSize,
    DataItem (..),
    JumpInstr (..),

    -- * Functions
    FuncDef (..),
    FuncParam (..),
    FuncArg (..),
    Abity (..),
    abityToBase,
    Block (..),
    fEntry,

    -- * Instructions
    Statement (..),
    Instr (..),
    VolatileInstr (..),
    ExtArg (..),
    toExtType,
    FloatArg (..),
    f2BaseType,
    IntArg (..),
    i2BaseType,
    IntCmpOp (..),
    FloatCmpOp (..),
    Phi (..),
    AllocSize (..),
    getSize,
  )
where

import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (fromJust)
import Data.Word (Word64)

-- TODO: Prefix all constructors

newtype UserIdent = UserIdent {userIdent :: String}
  deriving (Eq, Ord)

instance Show UserIdent where
  show (UserIdent s) = ':' : s

newtype LocalIdent = LocalIdent {localIdent :: String}
  deriving (Eq, Ord)

instance Show LocalIdent where
  show (LocalIdent s) = '%' : s

newtype BlockIdent = BlockIdent {blockIdent :: String}
  deriving (Eq, Ord)

instance Show BlockIdent where
  show (BlockIdent s) = '@' : s

newtype GlobalIdent = GlobalIdent {globalIdent :: String}
  deriving (Eq, Ord)

instance Show GlobalIdent where
  show (GlobalIdent s) = '$' : s

------------------------------------------------------------------------

data BaseType
  = Word
  | Long
  | Single
  | Double
  deriving (Show, Eq)

baseTypeByteSize :: BaseType -> Int
baseTypeByteSize Word = 4
baseTypeByteSize Long = 8
baseTypeByteSize Single = 4
baseTypeByteSize Double = 8

baseTypeBitSize :: BaseType -> Int
baseTypeBitSize ty = baseTypeByteSize ty * 8

data ExtType
  = Base BaseType
  | Byte
  | HalfWord
  deriving (Show, Eq)

extTypeByteSize :: ExtType -> Int
extTypeByteSize (Base b) = baseTypeByteSize b
extTypeByteSize Byte = 1
extTypeByteSize HalfWord = 2

extTypeBitSize :: ExtType -> Int
extTypeBitSize ty = extTypeByteSize ty * 8

data SubWordType
  = SignedByte
  | UnsignedByte
  | SignedHalf
  | UnsignedHalf
  deriving (Show, Eq)

data Abity
  = ABase BaseType
  | ASubWordType SubWordType
  | AUserDef UserIdent
  deriving (Show, Eq)

abityToBase :: Abity -> BaseType
-- Calls with a sub-word return type define a temporary of base type
-- w with its most significant bits unspecified.
abityToBase (ASubWordType _) = Word
-- When an aggregate type is used as argument type or return type, the
-- value respectively passed or returned needs to be a pointer to a
-- memory location holding the value.
abityToBase (AUserDef _) = Long
abityToBase (ABase ty) = ty

data Const
  = Number Word64
  | SFP Float
  | DFP Double
  | Global GlobalIdent
  deriving (Show, Eq)

data DynConst
  = Const Const
  | Thread GlobalIdent
  | Extern GlobalIdent
  | ExternThread GlobalIdent
  deriving (Show, Eq)

data Value
  = VConst DynConst
  | VLocal LocalIdent
  deriving (Show, Eq)

data Linkage
  = LExport
  | LThread
  | LSection String (Maybe String)
  deriving (Show, Eq)

data AllocSize
  = AllocWord
  | AllocLong
  | AllocLongLong
  deriving (Show, Eq)

getSize :: AllocSize -> Int
getSize AllocWord = 4
getSize AllocLong = 8
getSize AllocLongLong = 16

data TypeDef
  = TypeDef
  { aggName :: UserIdent,
    aggAlign :: Maybe Word64,
    aggType :: AggType
  }
  deriving (Show, Eq)

data SubType
  = SExtType ExtType
  | SUserDef UserIdent
  deriving (Show, Eq)

type Field = (SubType, Maybe Word64)

-- TODO: Type for tuple
data AggType
  = ARegular [Field]
  | AUnion [[Field]]
  | AOpaque Word64
  deriving (Show, Eq)

data DataDef
  = DataDef
  { linkage :: [Linkage],
    name :: GlobalIdent,
    align :: Maybe Word64,
    objs :: [DataObj]
  }
  deriving (Show, Eq)

dataSize :: DataDef -> Int
dataSize dataDef =
  sum $ map objSize (objs dataDef)

data DataObj
  = OItem ExtType [DataItem]
  | OZeroFill Word64
  deriving (Show, Eq)

objAlign :: DataObj -> Word64
objAlign (OZeroFill _) = 1 :: Word64
objAlign (OItem ty _) = fromIntegral $ extTypeByteSize ty

objSize :: DataObj -> Int
objSize (OZeroFill n) = fromIntegral n
objSize (OItem ty items) = extTypeByteSize ty * cnt items
  where
    cnt :: [DataItem] -> Int
    cnt [] = 0
    cnt ((DString s) : xs) = length s + cnt xs
    cnt (_ : xs) = 1 + cnt xs

data DataItem
  = DSymOff GlobalIdent Word64
  | DString String
  | DConst Const
  deriving (Show, Eq)

data FuncDef
  = FuncDef
  { fLinkage :: [Linkage],
    fName :: GlobalIdent,
    fStart :: BlockIdent,
    fAbity :: Maybe Abity,
    fParams :: [FuncParam],
    fBlock :: Map BlockIdent Block
  }
  deriving (Show, Eq)

fEntry :: FuncDef -> Block
fEntry func = fromJust $ Map.lookup (fStart func) (fBlock func)

data FuncParam
  = Regular Abity LocalIdent
  | Env LocalIdent
  | Variadic
  deriving (Show, Eq)

data FuncArg
  = ArgReg Abity Value
  | ArgEnv Value
  | ArgVar
  deriving (Show, Eq)

data JumpInstr
  = Jump BlockIdent
  | Jnz Value BlockIdent BlockIdent
  | Return (Maybe Value)
  | Halt
  deriving (Show, Eq)

data LoadType
  = LSubWord SubWordType
  | LBase BaseType
  deriving (Show, Eq)

-- TODO: Could/Should define this on ExtType instead.
loadByteSize :: LoadType -> Word64
loadByteSize (LSubWord UnsignedByte) = 1
loadByteSize (LSubWord SignedByte) = 1
loadByteSize (LSubWord SignedHalf) = 2
loadByteSize (LSubWord UnsignedHalf) = 2
loadByteSize (LBase Word) = 4
loadByteSize (LBase Long) = 8
loadByteSize (LBase Single) = 4
loadByteSize (LBase Double) = 8

data ExtArg
  = ExtSingle
  | ExtSubWord SubWordType
  | ExtSignedWord
  | ExtUnsignedWord
  deriving (Show, Eq)

toExtType :: ExtArg -> (Bool, ExtType)
toExtType (ExtSubWord SignedByte) = (True, Byte)
toExtType (ExtSubWord UnsignedByte) = (False, Byte)
toExtType (ExtSubWord SignedHalf) = (True, HalfWord)
toExtType (ExtSubWord UnsignedHalf) = (False, HalfWord)
toExtType ExtSignedWord = (True, Base Word)
toExtType ExtUnsignedWord = (False, Base Word)
toExtType ExtSingle = (True, Base Single)

data FloatArg = FDouble | FSingle
  deriving (Show, Eq)

f2BaseType :: FloatArg -> BaseType
f2BaseType FSingle = Single
f2BaseType FDouble = Double

data IntArg = IWord | ILong
  deriving (Show, Eq)

i2BaseType :: IntArg -> BaseType
i2BaseType IWord = Word
i2BaseType ILong = Long

-- TODO: Distinict types for floating point comparison?
data IntCmpOp
  = IEq
  | INe
  | ISle
  | ISlt
  | ISge
  | ISgt
  | IUle
  | IUlt
  | IUge
  | IUgt
  deriving (Show, Eq)

data FloatCmpOp
  = FEq
  | FNe
  | FLe
  | FLt
  | FGe
  | FGt
  | FOrd
  | FUnord
  deriving (Show, Eq)

data Instr
  = Add Value Value
  | Sub Value Value
  | Div Value Value
  | Mul Value Value
  | Neg Value
  | URem Value Value
  | Rem Value Value
  | UDiv Value Value
  | Or Value Value
  | Xor Value Value
  | And Value Value
  | Sar Value Value
  | Shr Value Value
  | Shl Value Value
  | Alloc AllocSize Value
  | Load LoadType Value
  | CompareInt IntArg IntCmpOp Value Value
  | CompareFloat FloatArg FloatCmpOp Value Value
  | Ext ExtArg Value
  | FloatToInt FloatArg Bool Value
  | IntToFloat IntArg Bool Value
  | TruncDouble Value
  | Cast Value
  | Copy Value
  | VAArg Value
  deriving (Show, Eq)

data VolatileInstr
  = Store ExtType Value Value
  | VAStart Value
  | Blit Value Value Word64
  | DBGLoc Word64 Word64 (Maybe Word64)
  deriving (Show, Eq)

data Statement
  = Assign LocalIdent BaseType Instr
  | Call (Maybe (LocalIdent, Abity)) Value [FuncArg]
  | Volatile VolatileInstr
  deriving (Show, Eq)

data Phi
  = Phi
  { pName :: LocalIdent,
    pType :: BaseType,
    pLabels :: Map BlockIdent Value
  }
  deriving (Show, Eq)

data Block
  = Block
  { label :: BlockIdent, -- TODO: Consider removing this (part of the Map)
    phi :: [Phi],
    stmt :: [Statement],
    term :: JumpInstr
  }
  deriving (Show, Eq)