-- The simplifier changes the shapes of closures that we expect.
{-# OPTIONS_GHC -O0 #-}
{-# LANGUAGE BangPatterns, MagicHash, UnboxedTuples #-}

import GHC.Exts.Heap

import Control.Concurrent.MVar
import Control.DeepSeq
import Control.Monad
import GHC.Exts
import GHC.Int
import GHC.IO
import GHC.IORef
import GHC.MVar
import GHC.Stack
import GHC.STRef
import GHC.Weak
import GHC.Word
import System.Environment
import System.Mem

exData :: (Int,Int)
exData = (1,2)

exItbl :: StgInfoTable
exItbl = StgInfoTable
    { entry = Nothing
    , ptrs = 0
    , nptrs = 0
    , tipe = toEnum 0
    , srtlen = 0
    , code = Nothing
    }

exConstrClosure :: Closure
exConstrClosure = ConstrClosure
    { info = exItbl{tipe=CONSTR_1_0, ptrs=1, nptrs=0}
    , ptrArgs = []
    , dataArgs = []
    , pkg = "base"
    , modl = "GHC.Base"
    , name = "Just"
    }

exFunClosure :: Closure
exFunClosure = FunClosure
    { info = exItbl{tipe=FUN_0_1, ptrs=0, nptrs=1}
    , ptrArgs = []
    , dataArgs = [0]
    }

exThunkClosure :: Closure
exThunkClosure = ThunkClosure
    { info = exItbl{tipe=THUNK}
    , ptrArgs = []
    , dataArgs = []
    }

exSelectClosure :: Closure
exSelectClosure = SelectorClosure
    { info = exItbl
    , selectee = asBox exData
    }

exPAPClosure :: Closure
exPAPClosure = PAPClosure
    { info = exItbl{tipe=PAP}
    , arity = 1
    , n_args = 1
    , fun = asBox id
    , payload = []
    }

exAPClosure :: Closure
exAPClosure = APClosure
    { info = exItbl{tipe=AP}
    , arity = 0
    , n_args = 0
    , fun = asBox id
    , payload = []
    }

exAPStackClosure :: Closure
exAPStackClosure = APStackClosure
    { info = exItbl{tipe=AP_STACK}
    , fun = asBox id
    , payload = []
    }

exIndClosure :: Closure
exIndClosure = IndClosure
    { info = exItbl{tipe=IND}
    , indirectee = asBox []
    }

exBCOClosure :: Closure
exBCOClosure = BCOClosure
    { info = exItbl{tipe=BCO, ptrs=4}
    , instrs = asBox []
    , literals = asBox []
    , bcoptrs = asBox []
    , arity = 0
    , size = 5
    , bitmap = []
    }

exBlackholeClosure :: Closure
exBlackholeClosure = BlackholeClosure
    { info = exItbl{tipe=BLACKHOLE}
    , indirectee = asBox []
    }

exArrWordsClosure :: Closure
exArrWordsClosure = ArrWordsClosure
    { info = exItbl{tipe=ARR_WORDS}
    , bytes = 0
    , arrWords = []
    }

exMutArrClosure :: Closure
exMutArrClosure = MutArrClosure
    { info = exItbl{tipe=MUT_ARR_PTRS_DIRTY}
    , mccPtrs = 0
    , mccSize = 0
    , mccPayload = []
    }

exMVarClosure :: Closure
exMVarClosure = MVarClosure
    { info = exItbl{tipe=MVAR_DIRTY}
    , queueHead = asBox []
    , queueTail = asBox []
    , value = asBox 0
    }

exMutVarClosure :: Closure
exMutVarClosure = MutVarClosure
    { info = exItbl{tipe=MUT_VAR_DIRTY}
    , var = asBox []
    }

exBlockingQClosure :: Closure
exBlockingQClosure = BlockingQueueClosure
    { info = exItbl{tipe=BLOCKING_QUEUE}
    , link = asBox []
    , blackHole = asBox []
    , owner = asBox []
    , queue = asBox []
    }

exWeakClosure :: Closure
exWeakClosure = WeakClosure
    { info = exItbl{tipe=WEAK}
    , cfinalizers = asBox []
    , key = asBox []
    , value = asBox []
    , finalizer = asBox []
    , link = asBox []
    }

exIntClosure :: Closure
exIntClosure = IntClosure
    { ptipe = PInt, intVal = 42 }

exWordClosure :: Closure
exWordClosure = WordClosure
    { ptipe = PWord, wordVal = 42 }

exInt64Closure :: Closure
exInt64Closure = Int64Closure
    { ptipe = PInt64, int64Val = 42 }

exWord64Closure :: Closure
exWord64Closure = Word64Closure
    { ptipe = PWord64, word64Val = 42 }

exAddrClosure :: Closure
exAddrClosure = AddrClosure
    { ptipe = PAddr, addrVal = 42 }

exFloatClosure :: Closure
exFloatClosure = FloatClosure
    { ptipe = PFloat, floatVal = 42.0 }

exDoubleClosure :: Closure
exDoubleClosure = DoubleClosure
    { ptipe = PDouble, doubleVal = 42.0 }

exOtherClosure :: Closure
exOtherClosure = OtherClosure
    { info = exItbl
    , hvalues = []
    , rawWords = []
    }

data A = A (Array# Int)
data MA = MA (MutableArray# RealWorld Int)
data BA = BA ByteArray#
data MBA = MBA (MutableByteArray# RealWorld)
data APC a = APC a

main :: IO ()
main = do

    --------------------------------------------
    -- Objects to inspect

    MA ma <- IO $ \s ->
        case newArray# 0# 0 s of
            (# s1, x #) -> (# s1, MA x #)
    A a <- IO $ \s ->
        case freezeArray# ma 0# 0# s of
            (# s1, x #) -> (# s1, A x #)
    MBA mba <- IO $ \s ->
        case newByteArray# 0# s of
            (# s1, x #) -> (# s1, MBA x #)
    BA ba <- IO $ \s ->
        case newByteArray# 0# s of
            (# s1, x #) ->
                case unsafeFreezeByteArray# x s1 of
                    (# s2, y #) -> (# s2, BA y #)
    bco <- IO $ \s ->
        newBCO# ba ba a 0# ba s
    APC apc <- IO $ \s ->
        case mkApUpd0# bco of
            (# x #) -> (# s, APC x #)

    --------------------------------------------
    -- Closures

    -- Constructor
    let !con = Just 1
    getClosureData con >>=
        assertClosuresEq exConstrClosure

    -- Function
    let !fun = \x -> x + 1
    getClosureData fun >>=
        assertClosuresEq exFunClosure

    -- Thunk
    let thunk = map (+2) [1,2,3]
    getClosureData thunk >>=
        assertClosuresEq exThunkClosure

    -- Selector
    -- FAILING: Getting THUNK not THUNK_SELECTOR
    -- let sel = case exData of (a,_) -> a
    -- getClosureData sel >>=
    --     assertClosuresEq exSelectClosure

    -- Partial application
    let !f = map (+2)
    getClosureData f >>=
        assertClosuresEq exPAPClosure

    -- Applied function
    getClosureData apc >>=
        assertClosuresEq exAPClosure

    -- Suspended thunk evaluation
    -- getClosureData (Just 1) >>=
    --     assertClosuresEq exAPStackClosure

    -- Indirection
    -- getClosureData (Just 1) >>=
    --     assertClosuresEq exIndClosure

    -- ByteCode object
    getClosureData bco >>=
        assertClosuresEq exBCOClosure

    -- Blackhole
    -- getClosureData (Just 1) >>=
    --     assertClosuresEq exBlackholeClosure

    -- Byte array
    getClosureData ba >>=
        assertClosuresEq exArrWordsClosure

    -- Mutable pointer array
    getClosureData ma >>=
        assertClosuresEq exMutArrClosure

    -- MVar
    (MVar v) <- newMVar 1
    getClosureData (unsafeCoerce# v) >>=
        assertClosuresEq exMVarClosure

    -- MutVar
    (IORef (STRef v)) <- newIORef 1
    getClosureData v >>=
        assertClosuresEq exMutVarClosure

    -- Blocking queue
    -- getClosureData (Just 1) >>=
    --    assertClosuresEq exBlockingQClosure

    -- Weak pointer
    Weak wk <- mkWeak (1 :: Int) (1 :: Int) Nothing

    getClosureData wk >>=
        assertClosuresEq exWeakClosure

    -----------------------------------------------------
    -- Unboxed unlifted types

    -- Primitive Int
    let (I# v) = 42
    getClosureData v >>=
        assertClosuresEq exIntClosure

    -- Primitive Word
    let (W# v) = 42
    getClosureData v >>=
        assertClosuresEq exWordClosure

    -- Primitive Int64
    -- FAILING: On 64-bit platforms, v is a regular Int
    -- let (I64# v) = 42
    -- getClosureData v >>=
    --     assertClosuresEq exInt64Closure

    -- Primitive Word64
    -- FAILING: On 64-bit platforms, v is a regular Word
    -- let (W64# v) = 42
    -- getClosureData v >>=
    --     assertClosuresEq exWord64Closure

    -- Primitive Addr
    let v = unsafeCoerce# 42# :: Addr#
    getClosureData v >>=
        assertClosuresEq exAddrClosure

    -- Primitive Float
    let (F# v) = 42.0
    getClosureData v >>=
        assertClosuresEq exFloatClosure

    -- Primitive Double
    let (D# v) = 42.0
    getClosureData v >>=
        assertClosuresEq exDoubleClosure

    ------------------------------------------------------
    -- Catch-all type

    -- Other
    -- getClosureData (Just 1) >>=
    --    assertClosuresEq exOtherClosure

    putStrLn "Done. No errors."


-- | Attempt to compare two closures
compareClosures :: Closure -> Closure -> Bool
compareClosures expected actual =
    -- Determine which fields to compare based
    -- upon expected closure type
    let funcs = case expected of
                    ConstrClosure{}         -> [ sEq (tipe . info)
                                               , sEq (ptrs . info)
                                               , sEq (nptrs . info)
                                               , sEq dataArgs
                                               , sEq name          ]
                    FunClosure{}            -> [ sEq (tipe . info)
                                               , sEq (ptrs . info)
                                               , sEq (nptrs . info)
                                               , sEq dataArgs      ]
                    ThunkClosure{}          -> [ sEq (tipe . info)
                                               , sEq (ptrs . info)
                                               , sEq (nptrs . info)
                                               , sEq dataArgs      ]
                    SelectorClosure{}       -> [ sEq (tipe . info) ]
                    PAPClosure{}            -> [ sEq (tipe . info)
                                               , sEq arity
                                               , sEq n_args        ]
                    APClosure{}             -> [ sEq (tipe . info)
                                               , sEq arity
                                               , sEq n_args        ]
                    APStackClosure{}        -> [ sEq (tipe . info) ]
                    IndClosure{}            -> [ sEq (tipe . info) ]
                    BCOClosure{}            -> [ sEq (tipe . info)
                                               , sEq arity
                                               , sEq bitmap        ]
                    BlackholeClosure{}      -> [ sEq (tipe . info) ]
                    ArrWordsClosure{}       -> [ sEq (tipe . info)
                                               , sEq bytes
                                               , sEq arrWords      ]
                    MutArrClosure{}         -> [ sEq (tipe . info)
                                               , sEq mccPtrs
                                               , sEq mccSize       ]
                    MVarClosure{}           -> [ sEq (tipe . info) ]
                    MutVarClosure{}         -> [ sEq (tipe . info) ]
                    BlockingQueueClosure{}  -> [ sEq (tipe . info) ]
                    WeakClosure{}           -> [ sEq (tipe . info) ]
                    IntClosure{}            -> [ sEq ptipe
                                               , sEq intVal    ]
                    WordClosure{}           -> [ sEq ptipe
                                               , sEq wordVal   ]
                    Int64Closure{}          -> [ sEq ptipe
                                               , sEq int64Val  ]
                    Word64Closure{}         -> [ sEq ptipe
                                               , sEq word64Val ]
                    AddrClosure{}           -> [ sEq ptipe
                                               , sEq addrVal   ]
                    FloatClosure{}          -> [ sEq ptipe
                                               , sEq floatVal  ]
                    DoubleClosure{}         -> [ sEq ptipe
                                               , sEq doubleVal ]
                    _ -> error $ "Don't know how to compare expected closure: "
                                 ++ show expected
    in compareWith funcs expected actual
  where
    -- Take a list of closure comparisons and check all
    compareWith :: [Closure -> Closure -> Bool] -> Closure -> Closure -> Bool
    compareWith funcs c1 c2 = all (\f -> f c1 c2) funcs

    -- Create a comparison function from a selector
    sEq :: Eq a => (Closure -> a) -> Closure -> Closure -> Bool
    sEq select c1 c2 = select c1 == select c2

-- | Assert two closures are equal, checking depending on closure type
assertClosuresEq :: HasCallStack => Closure -> Closure -> IO ()
assertClosuresEq _ c@UnsupportedClosure{} =
    fail $ unlines [ "Unsupported closure returned: " ++ show c
                   , ""
                   , prettyCallStack callStack
                   ]
assertClosuresEq expected actual =
    unless (compareClosures expected actual) $ fail $ unlines
        [ "assertClosuresEq: Closures do not match"
        , "Expected: " ++ show expected
        , "Actual:   " ++ show actual
        , ""
        , prettyCallStack callStack
        ]