{-# LANGUAGE RecordWildCards #-}

-- | Querying the type of LLVM expressions
module LLVM.AST.Typed (
  Typed(..),
  getElementType,
  indexTypeByConstants,
  indexTypeByOperands,
  extractValueType,
) where

import LLVM.Prelude

import Control.Monad.State (gets)
import qualified Data.Map.Lazy as Map
import qualified Data.Either as Either
import GHC.Stack

import LLVM.AST
import LLVM.AST.Global
import LLVM.AST.Type

import LLVM.IRBuilder.Module

import qualified LLVM.AST.Constant as C
import qualified LLVM.AST.Float as F

class Typed a where
  typeOf :: (HasCallStack, MonadModuleBuilder m) => a -> m (Either String Type)

instance Typed Operand where
  typeOf (LocalReference t _) = return $ Right t
  typeOf (ConstantOperand c)  = typeOf c
  typeOf _                    = return $ Right MetadataType

instance Typed CallableOperand where
  typeOf (Right op) = typeOf op
  typeOf (Left _) = return $ Left "typeOf inline assembler is not defined. (Malformed AST)"

instance Typed C.Constant where
  typeOf (C.Int bits _)       = return $ Right $ IntegerType bits
  typeOf (C.Float t)          = typeOf t
  typeOf (C.Null t)           = return $ Right t
  typeOf (C.AggregateZero t)  = return $ Right t
  typeOf (C.Struct {..}) = case structName of
                             Nothing -> do
                               mvtys <- mapM typeOf memberValues
                               case (all Either.isRight mvtys) of
                                 True -> return $ Right $ StructureType isPacked $ Either.rights mvtys
                                 False -> do
                                   let (Left s) = head $ filter Either.isLeft mvtys
                                   return $ Left $ "Could not deduce type for struct field: " ++ s
                             Just sn -> return $ Right $ NamedTypeReference sn
  typeOf (C.Array {..})  = return $ Right $ ArrayType (fromIntegral $ length memberValues) memberType
  typeOf (C.Vector {..}) = case memberValues of
                             []    -> return $ Left "Vectors of size zero are not allowed. (Malformed AST)"
                             (x:_) -> do
                               t <- typeOf x
                               case t of
                                 (Left _) -> return t
                                 (Right t') -> return $ Right $ VectorType (fromIntegral $ length memberValues) t'

  typeOf (C.Undef t)     = return $ Right t
  typeOf (C.BlockAddress {}) = return $ Right $ ptr i8
  typeOf (C.GlobalReference t _) = return $ Right t
  typeOf (C.Add {..})     = typeOf operand0
  typeOf (C.FAdd {..})    = typeOf operand0
  typeOf (C.FDiv {..})    = typeOf operand0
  typeOf (C.FRem {..})    = typeOf operand0
  typeOf (C.Sub {..})     = typeOf operand0
  typeOf (C.FSub {..})    = typeOf operand0
  typeOf (C.Mul {..})     = typeOf operand0
  typeOf (C.FMul {..})    = typeOf operand0
  typeOf (C.UDiv {..})    = typeOf operand0
  typeOf (C.SDiv {..})    = typeOf operand0
  typeOf (C.URem {..})    = typeOf operand0
  typeOf (C.SRem {..})    = typeOf operand0
  typeOf (C.Shl {..})     = typeOf operand0
  typeOf (C.LShr {..})    = typeOf operand0
  typeOf (C.AShr {..})    = typeOf operand0
  typeOf (C.And {..})     = typeOf operand0
  typeOf (C.Or  {..})     = typeOf operand0
  typeOf (C.Xor {..})     = typeOf operand0
  typeOf (C.GetElementPtr {..}) = do
    aty <- typeOf address
    case aty of
      (Left _) -> return aty
      (Right aty') -> indexTypeByConstants aty' indices
  typeOf (C.Trunc {..})    = return $ Right type'
  typeOf (C.ZExt {..})     = return $ Right type'
  typeOf (C.SExt {..})     = return $ Right type'
  typeOf (C.FPToUI {..})   = return $ Right type'
  typeOf (C.FPToSI {..})   = return $ Right type'
  typeOf (C.UIToFP {..})   = return $ Right type'
  typeOf (C.SIToFP {..})   = return $ Right type'
  typeOf (C.FPTrunc {..})  = return $ Right type'
  typeOf (C.FPExt {..})    = return $ Right type'
  typeOf (C.PtrToInt {..}) = return $ Right type'
  typeOf (C.IntToPtr {..}) = return $ Right type'
  typeOf (C.BitCast {..})  = return $ Right type'
  typeOf (C.ICmp {..})    = do
    t <- typeOf operand0
    case t of
      (Left _) -> return t
      (Right (VectorType n _)) -> return $ Right $ VectorType n i1
      (Right _) -> return $ Right i1
  typeOf (C.FCmp {..})    = do
    t <- typeOf operand0
    case t of
      (Left _) -> return t
      (Right (VectorType n _)) -> return $ Right $ VectorType n i1
      (Right _) -> return $ Right i1
  typeOf (C.Select {..})  = typeOf trueValue
  typeOf (C.ExtractElement {..})  = do
    t <- typeOf vector
    case t of
      (Left _) -> return t
      (Right (VectorType _ t')) -> return $ Right t'
      (Right _) -> return $ Left "The first operand of an extractelement instruction is a value of vector type. (Malformed AST)"
  typeOf (C.InsertElement {..})   = typeOf vector
  typeOf (C.ShuffleVector {..})   = do
    t0 <- typeOf operand0
    tm <- typeOf mask
    case (t0, tm) of
      (Right (VectorType _ t), Right (VectorType m _)) -> return $ Right $ VectorType m t
      _ -> return $ Left "The first operand of an shufflevector instruction is a value of vector type. (Malformed AST)"
  typeOf (C.ExtractValue {..})    = do
    t <- typeOf aggregate
    case t of
      (Left _) -> return t
      (Right t') -> extractValueType indices' t'
  typeOf (C.InsertValue {..})   = typeOf aggregate
  typeOf (C.TokenNone)          = return $ Right TokenType
  typeOf (C.AddrSpaceCast {..}) = return $ Right type'

-- | Index into a type using a list of 'Constant' values. Returns a pointer type whose referent is the indexed type, or an error message if indexing was not possible.
indexTypeByConstants :: (HasCallStack, MonadModuleBuilder m) => Type -> [C.Constant] -> m (Either String Type)
indexTypeByConstants ty [] = return $ Right $ ptr ty
indexTypeByConstants (PointerType ty _) (_:is) = indexTypeByConstants ty is
indexTypeByConstants (StructureType _ elTys) (C.Int 32 val:is) =
  indexTypeByConstants (elTys !! fromIntegral val) is
indexTypeByConstants (StructureType _ _) (i:_) =
  return $ Left $ "Indices into structures should be 32-bit integer constants. (Malformed AST): " ++ show i
indexTypeByConstants (VectorType _ elTy) (_:is) = indexTypeByConstants elTy is
indexTypeByConstants (ArrayType _ elTy) (_:is) = indexTypeByConstants elTy is
indexTypeByConstants (NamedTypeReference n) is = do
  mayTy <- liftModuleState (gets (Map.lookup n . builderTypeDefs))
  case mayTy of
    Nothing -> return $ Left $ "Couldn’t resolve typedef for: " ++ show n
    Just ty -> indexTypeByConstants ty is
indexTypeByConstants ty _ = return $ Left $ "Expecting aggregate type. (Malformed AST): " ++ show ty

-- | Index into a type using a list of 'Operand' values. Returns a pointer type whose referent is the indexed type, or an error message if indexing was not possible.
indexTypeByOperands :: (HasCallStack, MonadModuleBuilder m) => Type -> [Operand] -> m (Either String Type)
indexTypeByOperands ty [] = return $ Right $ ptr ty
indexTypeByOperands (PointerType ty _) (_:is) = indexTypeByOperands ty is
indexTypeByOperands (StructureType _ elTys) (ConstantOperand (C.Int 32 val):is) =
  indexTypeByOperands (elTys !! fromIntegral val) is
indexTypeByOperands (StructureType _ _) (i:_) =
  return $ Left $ "Indices into structures should be 32-bit integer constants. (Malformed AST): " ++ show i
indexTypeByOperands (VectorType _ elTy) (_:is) = indexTypeByOperands elTy is
indexTypeByOperands (ArrayType _ elTy) (_:is) = indexTypeByOperands elTy is
indexTypeByOperands (NamedTypeReference n) is = do
  mayTy <- liftModuleState (gets (Map.lookup n . builderTypeDefs))
  case mayTy of
    Nothing -> return $ Left $ "Couldn’t resolve typedef for: " ++ show n
    Just ty -> indexTypeByOperands ty is
indexTypeByOperands ty _ = return $ Left $ "Expecting aggregate type. (Malformed AST): " ++ show ty

getElementType :: Type -> Either String Type
getElementType (PointerType t _) = Right t
getElementType t = Left $ "Expecting pointer type. (Malformed AST): " ++ show t

extractValueType :: (HasCallStack, MonadModuleBuilder m) => [Word32] -> Type -> m (Either String Type)
extractValueType [] ty = return $ Right ty
extractValueType (i : is) (ArrayType numEls elTy)
  | fromIntegral i < numEls = extractValueType is elTy
  | fromIntegral i >= numEls = return $ Left $ "Expecting valid index into array type. (Malformed AST): " ++ show i
extractValueType (i : is) (StructureType _ elTys)
  | fromIntegral i < length elTys = extractValueType is (elTys !! fromIntegral i)
  | otherwise = return $ Left $ "Expecting valid index into structure type. (Malformed AST): " ++ show i
extractValueType _ ty = return $ Left $ "Expecting vector type. (Malformed AST): " ++ show ty

instance Typed F.SomeFloat where
  typeOf (F.Half _)          = return $ Right $ FloatingPointType HalfFP
  typeOf (F.Single _)        = return $ Right $ FloatingPointType FloatFP
  typeOf (F.Double _)        = return $ Right $ FloatingPointType DoubleFP
  typeOf (F.Quadruple _ _)   = return $ Right $ FloatingPointType FP128FP
  typeOf (F.X86_FP80 _ _)    = return $ Right $ FloatingPointType X86_FP80FP
  typeOf (F.PPC_FP128 _ _)   = return $ Right $ FloatingPointType PPC_FP128FP

instance Typed Global where
  typeOf (GlobalVariable {..}) = return $ Right $ type'
  typeOf (GlobalAlias {..})    = return $ Right $ type'
  typeOf (Function {..})       = do
    let (params, isVarArg) = parameters
    ptys <- mapM typeOf params
    case (all Either.isRight ptys) of
      True -> return $ Right $ FunctionType returnType (Either.rights ptys) isVarArg
      False -> do
        let (Left s) = head $ filter Either.isLeft ptys
        return $ Left $ "Could not deduce type for function parameter: " ++ s

instance Typed Parameter where
  typeOf (Parameter t _ _) = return $ Right t

instance Typed [Int32] where
  typeOf mask = return $ Right $ VectorType (fromIntegral $ length mask) i32