{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleInstances #-}

module LLVM.Typed (
  Typed(..),
  getElementType,
) where

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

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

-----
-- Reasoning about types
-----

class Typed a where
    typeOf :: a -> Type

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

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

instance Typed C.Constant where
    typeOf (C.Int bits _)  = IntegerType bits
    typeOf (C.Float float) = typeOf float
    typeOf (C.Null t)      = t
    typeOf (C.Struct {..}) = StructureType isPacked (map typeOf memberValues)
    typeOf (C.Array {..})  = ArrayType (fromIntegral $ length memberValues) memberType
    typeOf (C.Vector {..}) = VectorType (fromIntegral $ length memberValues) $
                                case memberValues of
                                    []    -> VoidType {- error "Vectors of size zero are not allowed" -}
                                    (x:_) -> typeOf x
    typeOf (C.Undef t)     = t
    typeOf (C.BlockAddress {..})   = ptr i8
    typeOf (C.GlobalReference t _) = 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 {..}) = getElementPtrType (typeOf address) indices
    typeOf (C.Trunc {..})   = type'
    typeOf (C.ZExt {..})    = type'
    typeOf (C.SExt {..})    = type'
    typeOf (C.FPToUI {..})  = type'
    typeOf (C.FPToSI {..})  = type'
    typeOf (C.UIToFP {..})  = type'
    typeOf (C.SIToFP {..})  = type'
    typeOf (C.FPTrunc {..}) = type'
    typeOf (C.FPExt {..})   = type'
    typeOf (C.PtrToInt {..}) = type'
    typeOf (C.IntToPtr {..}) = type'
    typeOf (C.BitCast {..})  = type'
    typeOf (C.ICmp {..})    = case (typeOf operand0) of
                                (VectorType n _) -> VectorType n i1
                                _ -> i1
    typeOf (C.FCmp {..})    = case (typeOf operand0) of
                                (VectorType n _) -> VectorType n i1
                                _ -> i1
    typeOf (C.Select {..})  = typeOf trueValue
    typeOf (C.ExtractElement {..})  = case typeOf vector of
                                        (VectorType _ t) -> t
                                        _ -> VoidType {- error "The first operand of an ‘extractelement‘ instruction is a value of vector type." -}
    typeOf (C.InsertElement {..})   = typeOf vector
    typeOf (C.ShuffleVector {..})   = case (typeOf operand0, typeOf mask) of
                                        (VectorType _ t, VectorType m _) -> VectorType m t
                                        _ -> VoidType {- error -}
    typeOf (C.ExtractValue {..})    = extractValueType (typeOf aggregate) indices
    typeOf (C.InsertValue {..})     = typeOf aggregate
    typeOf (C.TokenNone)          = TokenType
    typeOf (C.AddrSpaceCast {..}) = type'

getElementPtrType :: Type -> [C.Constant] -> Type
getElementPtrType ty [] = ptr ty
getElementPtrType (PointerType ty _) (_:is) = getElementPtrType ty is
getElementPtrType (StructureType _ elTys) (C.Int 32 val:is) =
  getElementPtrType (elTys !! fromIntegral val) is
getElementPtrType (VectorType _ elTy) (_:is) = getElementPtrType elTy is
getElementPtrType (ArrayType _ elTy) (_:is) = getElementPtrType elTy is

getElementType :: Type -> Type
getElementType (PointerType t _) = t
getElementType t = error $ "this should be a pointer type" ++ show t

extractValueType = error "extract"

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

instance Typed Global where
    typeOf (GlobalVariable {..}) = type'
    typeOf (GlobalAlias {..})    = type'
    typeOf (Function {..})       = let (params, isVarArg) = parameters
                                   in FunctionType returnType (map typeOf params) isVarArg
instance Typed Parameter where
    typeOf (Parameter t _ _) = t