packages feed

dx9d3d-0.1: DirectX9/D3D/Shader/Assembly.hs

module DirectX9.D3D.Shader.Assembly
where
import Data.Word    ( Word32 )
import Data.Bits    ( (.&.), (.|.), bit, testBit, shift )
import Control.Monad    ( when )
import System.IO    ( openBinaryFile, IOMode(ReadMode), hGetBuf )
import Foreign      ( allocaBytes, peekByteOff )

{-
 This code is based on directx9 sdk and
 http://msdn.microsoft.com/library/default.asp?url=/library/en-us/Display_r/hh/Display_r/usermodedisplaydriver_shader_cc577a5f-219b-47d0-916e-5a7f6eb75450.xml.asp

 Windows Codename Longhorn Display Driver Model Reference
  (Display and Print Devices: Windows Driver Kit)
 Direct3D Shader Codes
 
 So, the correct place might be WDK, but that's not confirmed.
-}

{-
 TODO
  * use continuations for generating stuff (toTokens)
  * use some parser lib or proper combinators for parsing (fromTokens)
-}

-- Helper classes, to unload namespace a bit

class Mask a where
    toMask      :: a -> Word32
    fromMask    :: Word32 -> Maybe a

class Token a where
    toToken     :: a -> Word32
    fromToken   :: Word32 -> Maybe a

class Tokens a where
    toTokens    :: a -> [Word32]
    fromTokens  :: [Word32] -> Maybe (a,[Word32])
{-
-- Not usefull, and causes error
instance (Token a) => Tokens [a] where
    toTokens = map toToken
    fromTokens = map fromToken
-}

instance (Tokens a) => Tokens [a] where
    toTokens = concatMap toTokens
    fromTokens l = case l of
        []  -> Nothing
        _   -> case fromTokens l of
            Nothing     -> Nothing
            Just (i,r)  -> build (i:[]) r
        where
            build :: (Tokens a) => [a] -> [Word32] -> Maybe ([a], [Word32])
            build res l = case l of
                [] -> return $ (reverse res, [])
                _:_  -> case fromTokens l of
                    Nothing     -> return (reverse res, l)
                    Just (i,r)  -> build (i:res) r

takeOne :: [Word32] -> Maybe (Word32, [Word32])
takeOne l = case l of
    []  -> Nothing
    i:x -> Just (i,x)                
{-
-- Requires undecideable instances
instance (Token a) => Tokens a where
    toTokens v = [toToken v]
    fromTokens v = case v of
        [] -> Nothing
        i:r -> case fromToken i of
            Nothing -> Nothing
            Just x  -> Just (x,r)
-}
---------------------------------------------------------------------------
-- Registers
---------------------------------------------------------------------------
-- Address is also Texture (for ps) and Output is also TexCoord (for vs >= 3_0)
data RegisterType
    = Temp | Input | Const | Address | Rasteriser | Attribute
    | Output | ConstInt | Color | Depth | Sampler | ConstBool
    | LoopReg | TempFloat16 | Misc | LabelReg | Predicate
    deriving (Show,Eq,Ord,Enum)

data Register = Register RegisterType Int deriving (Show,Eq,Ord)

instance Mask Register where
    toMask (Register rt i) = 
                (shift (val .&. (1+2+4)) 28)
            .|. (shift (val .&. (8+16))  8)
            .|. ((fromIntegral i') .&. 0x7ff)
        where (val, i') = case rt of
                Temp        -> (0,i)
                Input       -> (1,i)
                Const   -- There's actually 4 sets of constant registers for diffrent ranges
                    | i<2048    -> (2, i)
                    | i<4096    -> (11, i-2048)
                    | i<6144    -> (12, i-4096)
                    | otherwise -> (13, i-6144)
                Address     -> (3,i)  -- for ps, texture
                Rasteriser  -> (4,i)
                Attribute   -> (5,i)
                Output      -> (6,i) -- pre vs 3_0, texture
                ConstInt    -> (7,i)
                Color       -> (8,i)
                Depth       -> (9,i)
                Sampler     -> (10,i)
                ConstBool   -> (14,i)
                LoopReg     -> (15,i)
                TempFloat16 -> (16,i)
                Misc        -> (17,i)
                LabelReg    -> (18,i)
                Predicate   -> (19,i)
    fromMask mask = case mask' of
        0     -> Just $ Register Temp r'
        1     -> Just $ Register Input r'
        2     -> Just $ Register Const r'
        3     -> Just $ Register Address r'
        4     -> Just $ Register Rasteriser r'
        5     -> Just $ Register Attribute r'
        6     -> Just $ Register Output r'
        7     -> Just $ Register ConstInt r'
        8     -> Just $ Register Color r'
        9     -> Just $ Register Depth r'
        10    -> Just $ Register Sampler r'
        11    -> Just $ Register Const (r'+2048)
        12    -> Just $ Register Const (r'+4096)
        13    -> Just $ Register Const (r'+6144)
        14    -> Just $ Register ConstBool r'
        15    -> Just $ Register LoopReg r'
        16    -> Just $ Register TempFloat16 r'
        17    -> Just $ Register Misc r'
        18    -> Just $ Register LabelReg r'
        19    -> Just $ Register Predicate r'
        _       -> Nothing
        where
            mask' = ((shift mask (-28)) .&. (1+2+4)) .|. ((shift mask (-8)) .&. (8+16))
            r' = fromIntegral $ mask .&. 0x7ff



---------------------------------------------------------------------------
-- SourceModifier
---------------------------------------------------------------------------

-- Boolean variable means negation
data SourceModifier
    = ModNone   Bool
    | ModBias   Bool
    | ModSign   Bool
    | ModComplement
    | ModX2     Bool
    | ModDz
    | ModDw
    | ModAbs    Bool
    | ModNot
    deriving (Show,Eq,Ord)

instance Mask SourceModifier where
    toMask a = case a of
        ModNone n   -> prep 0x0 n
        ModBias n   -> prep 0x2 n
        ModSign n   -> prep 0x4 n
        ModComplement -> prep 0x6 False
        ModX2 n     -> prep 0x7 n
        ModDz       -> prep 0x9 False
        ModDw       -> prep 0xa False
        ModAbs n    -> prep 0xb n
        ModNot      -> prep 0xd False
        where prep x n = flip shift 24 $ x + if n then 1 else 0
    fromMask m = case ((shift m (-24)) .&. 0xf ) of
        0x0 -> Just $ ModNone False
        0x1 -> Just $ ModNone True
        0x2 -> Just $ ModBias False
        0x3 -> Just $ ModBias True
        0x4 -> Just $ ModSign False
        0x5 -> Just $ ModSign True
        0x6 -> Just $ ModComplement
        0x7 -> Just $ ModX2 False
        0x8 -> Just $ ModX2 True
        0x9 -> Just $ ModDz
        0xa -> Just $ ModDw
        0xb -> Just $ ModAbs False
        0xc -> Just $ ModAbs True
        0xd -> Just $ ModNot
        _   -> Nothing

---------------------------------------------------------------------------
-- Swizzle
---------------------------------------------------------------------------
data SwizzleSource
    = ComponentX | ComponentY | ComponentZ | ComponentW
    deriving (Show,Eq,Ord)

data Swizzle
    = Swizzle SwizzleSource SwizzleSource SwizzleSource SwizzleSource
    deriving (Show,Eq,Ord)

instance Mask Swizzle where
    toMask (Swizzle x y z w) =
            (prep 16 x) .|. (prep 18 y)
        .|. (prep 20 z) .|. (prep 22 w)
        where prep n s = flip shift n $ case s of
                ComponentX  -> 0x0
                ComponentY  -> 0x1
                ComponentZ  -> 0x2
                ComponentW  -> 0x3

    fromMask m = Just $ Swizzle (prep 16) (prep 18) (prep 20) (prep 22)
        where
            b n = (shift m (-n) .&. 0x3)
            prep n
                | b n == 0x0    = ComponentX
                | b n == 0x1    = ComponentY
                | b n == 0x2    = ComponentZ
                | b n == 0x3    = ComponentW
                | otherwise     = ComponentX  -- to shut ghc up on pattern

---------------------------------------------------------------------------
-- Source
---------------------------------------------------------------------------
data SourceReg = SourceReg Register SourceModifier Swizzle
    deriving (Show,Eq,Ord)

data Source = Source SourceReg (Maybe SourceReg)
    deriving (Show,Eq,Ord)

instance Mask SourceReg where
    toMask (SourceReg r sm s) =
        toMask r .|. toMask sm .|. toMask s .|. bit 31

    fromMask m = do
            r <- fromMask m
            sm <- fromMask m
            s <- fromMask m
            return $ SourceReg r sm s

instance Token SourceReg where
    toToken = toMask
    fromToken = fromMask

instance Tokens Source where
    toTokens (Source r ra) = case ra of
        Nothing -> [build r False]
        Just x  -> [build r True, build x False]
        where
            build r ra = toMask r .|. (if ra then bit 13 else 0)
    fromTokens l = do
        (x,l') <- takeOne l
        r <- fromMask x
        if testBit x 13
            then do
                (x', l'') <- takeOne l'
                ra <- fromMask x'
                return $ (Source r (Just ra), l'')
            else
                return $ (Source r Nothing, l')

        
---------------------------------------------------------------------------
-- DestinationModifier
---------------------------------------------------------------------------
data DestinationModifiers = DestinationModifiers Bool Bool Bool
    deriving (Show,Eq,Ord)

instance Mask DestinationModifiers where
    toMask (DestinationModifiers s pp c) =
        flip shift 20
            $   (if s  then 0x1 else 0)
            .|. (if pp then 0x2 else 0)
            .|. (if c  then 0x4 else 0)
    fromMask m = Just $ DestinationModifiers
        (testBit m 20) (testBit m 21) (testBit m 22)

---------------------------------------------------------------------------
-- Destination
---------------------------------------------------------------------------
data WriteMask = WriteMask Bool Bool Bool Bool
    deriving (Show,Eq,Ord)
instance Mask WriteMask where
    toMask (WriteMask x y z w) = 
            (if x then bit 16 else 0)
        .|. (if y then bit 17 else 0)
        .|. (if z then bit 18 else 0)
        .|. (if w then bit 19 else 0)
    fromMask m = Just $ WriteMask
        (testBit m 16) (testBit m 17) (testBit m 18) (testBit m 19)

data DestinationReg = DestinationReg Register WriteMask DestinationModifiers
    deriving (Show,Eq,Ord)

data Destination = Destination DestinationReg (Maybe DestinationReg)
    deriving (Show,Eq,Ord)

instance Mask DestinationReg where
    toMask (DestinationReg r wm dm) =
        bit 31 .|. toMask wm .|. toMask dm .|. toMask r

    fromMask m = do
            r <- fromMask m
            wm <- fromMask m
            dm <- fromMask m
            return $ DestinationReg r wm dm

instance Token DestinationReg where
    toToken = toMask
    fromToken = fromMask

instance Tokens Destination where
    toTokens (Destination r ra) = case ra of
        Nothing -> [build r False]
        Just x  -> [build r True, build x False]
        where build r ra = toMask r .|. (if ra then bit 13 else 0)
    fromTokens l = do
        (x,l') <- takeOne l
        r <- fromMask x
        if testBit x 13
            then do
                (x', l'') <- takeOne l'
                ra <- fromMask x'
                return $ (Destination r (Just ra), l'')
            else
                return $ (Destination r Nothing, l')

---------------------------------------------------------------------------
-- OpCode
---------------------------------------------------------------------------
data OpCode
    = OpNOP | OpMOV | OpADD | OpSUB | OpMAD | OpMUL | OpRCP | OpRSQ | OpDP3 | OpDP4
    | OpMIN | OpMAX | OpSLT | OpSGE | OpEXP | OpLOG | OpLIT | OpDST | OpLRP | OpFRC
    | OpM4x4 | OpM4x3 | OpM3x4 | OpM3x3 | OpM3x2 | OpCALL | OpCALLNZ | OpLOOP | OpRET
    | OpENDLOOP | OpLABEL | OpDCL | OpPOW | OpCRS | OpSGN | OpABS | OpNRM | OpSINCOS   
    | OpREP | OpENDREP | OpIF | OpIFC | OpELSE | OpENDIF | OpBREAK  | OpBREAKC  | OpMOVA     
    | OpDEFB | OpDEFI     

    | OpTEXCOORD | OpTEXKILL | OpTEX | OpTEXBEM | OpTEXBEML | OpTEXREG2AR | OpTEXREG2GB   
    | OpTEXM3x2PAD | OpTEXM3x2TEX | OpTEXM3x3PAD | OpTEXM3x3TEX | OpRESERVED0 | OpTEXM3x3SPEC 
    | OpTEXM3x3VSPEC | OpEXPP | OpLOGP | OpCND | OpDEF | OpTEXREG2RGB | OpTEXDP3TEX   
    | OpTEXM3x2DEPTH | OpTEXDP3 | OpTEXM3x3 | OpTEXDEPTH | OpCMP | OpBEM | OpDP2ADD      
    | OpDSX | OpDSY | OpTEXLDD | OpSETP | OpTEXLDL | OpBREAKP

    | OpPHASE | OpCOMMENT | OpEND
    deriving (Show, Eq, Ord, Enum)

opToV :: OpCode -> Word32
opToV o = case o of
    OpPHASE     -> 0xfffd
    OpCOMMENT   -> 0xfffe
    OpEND       -> 0xffff
    _
        | o<OpTEXCOORD  -> fromIntegral $ fromEnum o
        | otherwise     -> fromIntegral $ fromEnum o + 15

vToOp :: Word32 -> Maybe OpCode
vToOp v
    | v==0xfffd     = Just OpPHASE
    | v==0xfffe     = Just OpCOMMENT
    | v==0xffff     = Just OpEND
    | v<64          = Just $ toEnum $ fromIntegral v
    | v<97          = Just $ toEnum $ fromIntegral $ v - 15
    | otherwise     = Nothing

---------------------------------------------------------------------------
-- Instruction
---------------------------------------------------------------------------
-- This cannot handle dcl_xyz or version opcodes
data Instruction = Instruction OpCode Word32 Bool Bool (Maybe Destination) [Source]
    deriving (Show,Eq,Ord)

buildInstructionToken :: Word32 -> Word32 -> Word32 -> Bool ->  Bool -> Word32
buildInstructionToken code control size pred coissue =
    sm 16 0xff control . sm 24 0xf size . smc 28 pred . smc 30 coissue $ (code .&. 0xffff)
    where
        sm b m v n = (shift (m .&. v) b) .|. n
        smc b v n
            | v         = bit b .|. n
            | otherwise = n

fromInstructionToken :: Word32 -> (Word32, Word32, Word32, Bool, Bool)
fromInstructionToken m = (code, control, size, pred, coissue)
    where
        code = m .&. 0xffff
        control = (shift m (-16)) .&. 0xff
        size = (shift m (-24)) .&. 0xf
        pred = testBit m 28
        coissue = testBit m 30

instance Tokens Instruction where
    toTokens (Instruction code ctrl pred co d s) =
        (buildInstructionToken (opToV code) ctrl len pred co):ts
        where
            dest = maybe [] toTokens d
            sour = concatMap toTokens s
            ts = dest++sour
            len = fromIntegral $ length ts
        {- params (buildInstructionToken (opToV code) ctrl size pred co)
        where
            size = maybe 0 (const 1) d + fromIntegral (length s)
            params n = case d of
                Just x -> n:toToken x:map toToken s
                Nothing -> n:map toToken s -}

    fromTokens ts = case ts of
        []      -> Nothing
        i:ts    -> let
            (code,ctrl,size,pred,co) = fromInstructionToken i
            in if size==0
                then do
                    op <- vToOp code
                    return (Instruction op ctrl pred co Nothing [], ts)
                else do
                    op          <- vToOp code
                    (d,ts')     <- fromTokens ts
                    (ss,ts'')   <- many (size-1) ts'
                    return (Instruction op ctrl pred co (Just d) ss, ts'')
        where
            many = many' []
            many' res i l
                | i==0      = Just (reverse res, l)
                | otherwise = do
                    (t,l') <- fromTokens l
                    many' (t:res) (i-1) l'

---------------------------------------------------------------------------
-- Version
---------------------------------------------------------------------------
data Version = PsVersion Int Int | VsVersion Int Int
    deriving (Show,Eq,Ord)

instance Token Version where
    toToken v = case v of
        PsVersion maj min     -> b 0xffff0000 maj min
        VsVersion maj min    -> b 0xfffe0000 maj min
        where b t ma mi = t .|. (shift (fromIntegral ma) 8) .|. (fromIntegral mi)
    fromToken i= let
            t = i .&. 0xffff0000
            min = fromIntegral $ i .&. 0xff
            maj = fromIntegral $ shift (i .&. 0xff00) (-8)
            in if       t==0xffff0000   then Just $ PsVersion maj min
                else if t==0xfffe0000   then Just $ VsVersion maj min
                                        else Nothing



---------------------------------------------------------------------------
-- DeclUsage
---------------------------------------------------------------------------
data DeclUsageType
    = DuPosition | DuBlendWeight | DuBlenIndices | DuNormal | DuPSize
    | DuTexCoord | DuTangent | DuBiNormal | DuTessFactor | DuPositionT
    | DuColor | DuFog | DuDepth | DuSample
    deriving (Show, Eq, Ord, Enum)

data DeclUsage = DeclUsage DeclUsageType Int
    deriving (Show, Eq, Ord)

dutToV :: DeclUsageType -> Word32
dutToV = fromIntegral . fromEnum

vToDut :: Word32 -> Maybe DeclUsageType
vToDut v
    | v<14      = Just $ toEnum $ fromIntegral v
    | otherwise = Nothing


instance Token DeclUsage where
    toToken (DeclUsage decl usage) =
        (shift ((fromIntegral usage) .&. 0xf) 16) .|. bit 31 .|. (dutToV decl)
    fromToken w = do
        let
            t = w .&. 0x1f
            u = (shift w (-16)) .&. 0xf
        t' <- vToDut t
        return (DeclUsage t' $ fromIntegral u)

---------------------------------------------------------------------------
-- DeclarePS
---------------------------------------------------------------------------
data SamplerType = St2D | StCube | StVolume deriving (Show, Eq, Ord)
data DeclareTypePS
    = PsSampler SamplerType -- regnum used, regtype must be sampler
    | PsInput
    | PsTexture Int
    | PsColor Int
    deriving (Show, Eq, Ord)

data DeclarePs = DeclarePs DeclareTypePS DestinationReg
    deriving (Show, Eq, Ord)

instance Tokens DeclarePs where
    toTokens (DeclarePs t d) = case t of
        PsSampler st    -> let
            texType = (bit 31 .|.) $ flip shift 27 $ case st of
                St2D -> 2; StCube -> 3; StVolume -> 4;
            ins = buildInstructionToken 31 0 2 False False
            in [ins, texType, toToken d]
        PsInput         -> [buildInstructionToken 31 0 2 False False, bit 31, toToken d]
        PsTexture u     -> [buildInstructionToken 31 0 2 False False, toToken (DeclUsage DuTexCoord u), toToken d]
        PsColor u       -> [buildInstructionToken 31 0 2 False False, toToken (DeclUsage DuColor u), toToken d]
    fromTokens l = case l of
        i:t:d:r -> do
            let
                (code,_,size,_,co) = fromInstructionToken i
                usaget = (shift t (-16)) .&. 0xf
                -- usage = t .&. 0x1f
                sampler = (shift t (-27)) .&. 0xf
            when (code/=31 || size/=2 || co) Nothing
            d' <- fromToken d
            if sampler/=0
                then do
                    sampler' <- case sampler of 2 -> return St2D; 3 -> return StCube; 4 -> return StVolume; _->Nothing
                    return (DeclarePs (PsSampler sampler') d',r)
                else if usaget/=0 then do
                    du <- fromToken t
                    case du of
                        DeclUsage DuTexCoord u  -> return (DeclarePs (PsTexture u) d',r)
                        DeclUsage DuColor u     -> return (DeclarePs (PsColor u) d',r)
                        _   -> Nothing
                else return (DeclarePs PsInput d',r)
        _       -> Nothing

---------------------------------------------------------------------------
-- DeclareVS
---------------------------------------------------------------------------
data DeclareVs = DeclareVs DeclUsage DestinationReg
    deriving (Show, Eq, Ord)

instance Tokens DeclareVs where
    toTokens (DeclareVs u d) = [buildInstructionToken 31 0 2 False False, toToken u, toToken d]
    fromTokens l = case l of
        i:t:d:r -> do
            let
                (code,_,size,_,co) = fromInstructionToken i
            when (code/=31 || size/=2 || co) Nothing
            t' <- fromToken t
            d' <- fromToken d
            return (DeclareVs t' d',r)
        _       -> Nothing



---------------------------------------------------------------------------
-- ShaderAssembly
---------------------------------------------------------------------------

data ShaderAssembly
    = PixelShader Version [DeclarePs] [Instruction]
    | VertexShader Version [DeclareVs] [Instruction]
    deriving (Show)

instance Tokens ShaderAssembly where
    toTokens a = case a of
        PixelShader v d i   -> toToken v:toTokens d ++ toTokens i
        VertexShader v d i  -> toToken v:toTokens d ++ toTokens i
    fromTokens l = do
        (v,l') <- multi l
        case v of
            PsVersion _ _ -> do
                (d,l'') <- fromTokens l'
                (i,l''') <- fromTokens l''
                return $ (PixelShader v d i, l''')
            VsVersion _ _ -> do
                (d,l'') <- fromTokens l'
                (i,l''') <- fromTokens l''
                return $ (VertexShader v d i, l''')
        where
            multi l = case l of
                []  -> Nothing
                i:x -> do
                    t <- fromToken i
                    return (t,x)

disasm :: [Word32] -> Maybe ShaderAssembly
disasm l = fromTokens l >>= return.fst

---------------------------------------------------------------------------
-- Helpers
---------------------------------------------------------------------------
readWord32File :: String -> IO [Word32]
readWord32File f = do
    h <- openBinaryFile f ReadMode
    res <- allocaBytes 1024 $ readAll [] h 1024
    return $ concat res
    where
        readAll ret h s ptr = do
            s' <- hGetBuf h ptr s
            res <- flip mapM [0..s' `div` 4-1] $ \i -> do
                peekByteOff ptr (i*4)
            if s>s'
                then return $ (res:ret)
                else readAll (res:ret) h s ptr

readShaderAssembly :: String -> IO (Maybe ShaderAssembly)
readShaderAssembly f = readWord32File f >>= return.disasm