{-# LANGUAGE CPP             #-}
{-# LANGUAGE DataKinds       #-}
{-# LANGUAGE ImplicitParams  #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeOperators   #-}

{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

-- | Check properties of Ivory programs using random inputs.
--
-- Example usage:
--
-- > [ivory|
-- > struct foo
-- >   { foo_a :: Stored IFloat
-- >   ; foo_b :: Stored Uint8
-- >   }
-- > |]
-- >
-- > -- Function we want to generate inputs for.
-- > func :: Def ('[ Uint8
-- >               , Ref s (Array 3 (Stored Uint8))
-- >               , Ref s (Struct "foo")
-- >               ] :-> ())
-- > func = proc "func" $ \u arr str ->
-- >   ensures_ (checkStored (arr ! 0) (\r -> r >? u)) $
-- >   body $
-- >   arrayMap $ \ix -> do
-- >     a <- deref (arr ! ix)
-- >     b <- deref (str ~> foo_b)
-- >     store (arr ! ix) (a + b + u)
-- >
-- > -- Module containing our function
-- > cmodule :: Module
-- > cmodule = package "module" $ do
-- >   defStruct (Proxy :: Proxy "foo")
-- >   incl func
-- >
-- > -- Running @mkTest@ will produce a C program in @<pwd>/test@ that will check
-- > -- @func@'s contract on 10 random inputs.
-- > mkTest :: IO ()
-- > mkTest = check 10 [] cmodule (contract func)

module Ivory.QuickCheck (check, checkWith, contract) where

import           Prelude                         ()
import           Prelude.Compat

import           Control.Monad                   (forM, replicateM)
import           Data.IORef                      (IORef, newIORef, readIORef,
                                                  writeIORef)
import           Data.List                       (find, transpose)
import           Ivory.Compile.C.CmdlineFrontend
import qualified Ivory.Eval                      as E
import           Ivory.Language
import           Ivory.Language.Proc
import qualified Ivory.Language.Syntax           as I
import           System.IO.Unsafe                (unsafeInterleaveIO)

import qualified Test.QuickCheck.Arbitrary       as A
import qualified Test.QuickCheck.Gen             as G

import           Data.Int
import           Data.Word

-- XXX: DEBUG
-- import Debug.Trace

-- | Generate a random C program to check that the property holds. The
-- generated program will be placed in the @test@ subdirectory.
checkWith :: Int                  -- ^ The number of inputs to generate.
          -> Opts                 -- ^ Options to pass to the Ivory compiler.
          -> [Module]             -- ^ Modules we need to have in scope.
          -> Module               -- ^ The defining module.
          -> Def (args ':-> IBool) -- ^ The property to check.
          -> IO ()
checkWith n opts deps m prop@(DefProc p) = do
  inputs <- sampleProc m p n
  let main = DefProc I.Proc { I.procSym = "main"
                            , I.procRetTy = I.TyInt I.Int32
                            , I.procArgs = []
                            , I.procBody = concat inputs ++ [I.Return $ I.Typed (I.TyInt I.Int32) (I.ExpLit (I.LitInteger 0))]
                            , I.procRequires = []
                            , I.procEnsures = []
                            }
  let test = package (I.modName m ++ "__test") $ do
               depend m
               incl prop
               incl main
  runCompiler ([m, test]++deps) [] opts -- initialOpts { outDir = Just "test" }
checkWith _ _ _ _ _ = error "I can only check normal Ivory procs!"

-- | Generate a random C program to check that the property holds. The
-- generated program will be placed in the @test@ subdirectory.
check :: Int                  -- ^ The number of inputs to generate.
      -> [Module]             -- ^ Modules we need to have in scope.
      -> Module               -- ^ The defining module.
      -> Def (args ':-> IBool) -- ^ The property to check.
      -> IO ()
check n deps m p = checkWith n (initialOpts { outDir = Just "test"}) deps m p

-- | Make a @check@able property from an arbitrary Ivory procedure. The
-- property will simply check that the contracts are satisfied.
contract :: Def (args ':-> res) -> Def (args ':-> IBool)
contract (DefProc (I.Proc {..}))
  = DefProc I.Proc
      { I.procSym = procSym ++ "__contract_check"
      , I.procRetTy = I.TyBool
      , I.procArgs = procArgs
      , I.procBody = [ I.Call procRetTy Nothing (I.NameSym procSym)
                       [ I.Typed t (I.ExpVar v) | I.Typed t v <- procArgs ]
                     , I.Return (I.Typed I.TyBool (I.ExpLit (I.LitBool True)))
                     ]
      , I.procRequires = procRequires
      , I.procEnsures = []
      }
contract _ = error "I can only check contracts of normal Ivory procs!"

mkUnique :: (?counter :: IORef Integer) => IO Integer
mkUnique = do
  i <- readIORef ?counter
  writeIORef ?counter (i+1)
  return i

sampleProc :: Module -> I.Proc -> Int -> IO [I.Block]
sampleProc m@(I.Module {..}) p@(I.Proc {..}) n
  = do c <- newIORef 0
       let ?counter = c
       allInits <- fmap transpose $ forM procArgs $ \ (I.Typed t _) ->
         sampleType m t
       allAreas <- fmap transpose $ forM (getVisible modAreas) $ \ I.Area {..} -> do
         inits <- sampleType m areaType
         lazyMapIO (\ (var, blck) -> do
                     let store = case areaType of
                           I.TyArr _ _  -> I.RefCopy
                           I.TyStruct _ -> I.RefCopy
                           I.TyRef _    -> I.RefCopy -- XXX: this shouldn't actually appear
                           _            -> I.Store
                     return $ blck ++ [ store areaType (I.ExpAddrOfGlobal areaSym) (I.ExpVar var)])
                   inits
       --XXX: Refactor!
       let validInits =
             [ (inits, areas) | (inits, areas) <- zipLonger allInits allAreas
             , let (vars, blcks) = unzip inits
             , let asgnv = [ I.Assign ty arg (I.ExpVar var)
                           | (I.Typed ty arg, var) <- zip procArgs vars
                           ]
             , E.runEval (E.openModule m (do
                 E.evalBlock (concat blcks ++ asgnv ++ concat areas)
                 E.evalRequires procRequires))
               == Right True
             ]
       forM (take n validInits) $ \ (args, areas) -> do
         let (vars, inits) = unzip args
         chk <- mkCheck p vars
         return (concat inits ++ concat areas ++ chk)

getVisible :: I.Visible a -> [a]
getVisible xs = I.public xs ++ I.private xs

mkCheck :: (?counter :: IORef Integer)
        => I.Proc -> [I.Var] -> IO I.Block
mkCheck (I.Proc {..}) args = do
  n <- mkUnique
  let b = mkVar n
  let c = [ I.Call I.TyBool (Just b) (I.NameSym procSym)
            [ I.Typed t (I.ExpVar v) | (I.Typed t _, v) <- zip procArgs args ]
          , I.Assert (I.ExpVar b) ]
  return c

sampleType :: (?counter :: IORef Integer)
           => Module -> I.Type -> IO [(I.Var, I.Block)]
sampleType m t = case t of
  I.TyInt sz
    -> lazyMapIO (mkLocal t) =<< sampleInt sz
  I.TyWord sz
    -> lazyMapIO (mkLocal t) =<< sampleWord sz
  I.TyIndex i
    -> lazyMapIO (mkLocal t) =<< sampleIndex i
  I.TyBool
    -> lazyMapIO (mkLocal t) =<< sampleBool
  I.TyChar
    -> lazyMapIO (mkLocal t) =<< sampleChar
  I.TyFloat
    -> lazyMapIO (mkLocal t) =<< sampleFloat
  I.TyDouble
    -> lazyMapIO (mkLocal t) =<< sampleDouble
  I.TyRef ty
    -> lazyMapIO (mkRef ty) =<< sampleType m ty
  I.TyConstRef ty
    -> lazyMapIO (mkRef ty) =<< sampleType m ty
  I.TyArr len ty
    -> sampleArray m len ty
  I.TyStruct ty
    -> sampleStruct m ty
  I.TyProc _ _    -> err
  I.TyVoid        -> err
  I.TyPtr _       -> err
  I.TyConstPtr _  -> err
  I.TyCArray _    -> err
  I.TyOpaque      -> err
  where
  err = error $ "I don't know how to make values of type '" ++ show t ++ "'!"

mkLocal :: (?counter :: IORef Integer) => I.Type -> I.Init -> IO (I.Var, I.Block)
mkLocal ty init = do
  n <- mkUnique
  let v = mkVar n
  return (v, [I.Local ty v init])

mkRef :: (?counter :: IORef Integer) => I.Type -> (I.Var, I.Block)
      -> IO (I.Var, I.Block)
mkRef ty (v, init) = do
  n <- mkUnique
  let r = I.VarName ("ref" ++ show n)
  return (r, init ++ [ I.AllocRef ty r (I.NameVar v) ])

mkVar :: Integer -> I.Var
mkVar n = I.VarName ("var" ++ show n)

mk :: G.Gen a -> IO [a]
mk g = G.generate (G.infiniteListOf g)

sampleInt :: I.IntSize -> IO [I.Init]
sampleInt sz = do
  xs <- gen
  return [ I.InitExpr (I.TyInt sz) (I.ExpLit (I.LitInteger x)) | x <- xs ]
  where
  gen = case sz of
    I.Int8  -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Int8)
    I.Int16 -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Int16)
    I.Int32 -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Int32)
    I.Int64 -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Int64)

sampleWord :: I.WordSize -> IO [I.Init]
sampleWord sz = do
  xs <- gen
  return [ I.InitExpr (I.TyWord sz) (I.ExpLit (I.LitInteger x))
         | x <- xs
         ]
  where
  gen = case sz of
    I.Word8  -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Word8)
    I.Word16 -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Word16)
    I.Word32 -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Word32)
    I.Word64 -> fmap fromIntegral <$> mk (A.arbitrary :: G.Gen Word64)

sampleIndex :: Integer -> IO [I.Init]
sampleIndex ix = do
  xs <- mk (A.arbitrary :: G.Gen Word64)
  return [ I.InitExpr (I.TyIndex ix)
           (I.ExpLit (I.LitInteger (fromIntegral x `mod` ix)))
         | x <- xs
         ]

sampleBool :: IO [I.Init]
sampleBool = do
  bs <- mk A.arbitrary
  return [I.InitExpr I.TyBool (I.ExpLit (I.LitBool b)) | b <- bs]

sampleChar :: IO [I.Init]
sampleChar = do
  cs <- mk A.arbitrary
  return [I.InitExpr I.TyChar (I.ExpLit (I.LitChar c)) | c <- cs]

sampleFloat :: IO [I.Init]
sampleFloat = do
  cs <- mk A.arbitrary
  return [I.InitExpr I.TyFloat (I.ExpLit (I.LitFloat c)) | c <- cs]

sampleDouble :: IO [I.Init]
sampleDouble = do
  cs <- mk A.arbitrary
  return [I.InitExpr I.TyDouble (I.ExpLit (I.LitDouble c)) | c <- cs]

sampleStruct :: (?counter :: IORef Integer)
             => I.Module -> String -> IO [(I.Var, I.Block)]
sampleStruct m@(I.Module {..}) ty
  = case find (\s -> ty == I.structName s) structs of
    Just (I.Struct _ fields) -> repeatIO $ do
      (vars, blcks) <- fmap unzip $ forM fields $ \ (I.Typed t _) ->
        head <$> sampleType m t
      let init = zipWith (\ (I.Typed t f) v -> (f, I.InitExpr t (I.ExpVar v)))
                 fields vars
      (v, blck) <- mkLocal (I.TyStruct ty) (I.InitStruct init)
      return (v, concat blcks ++ blck)
    _ -> error ("I don't know how to construct a '" ++ ty ++ "'!")
  where
  structs = I.public modStructs ++ I.private modStructs

sampleArray :: (?counter :: IORef Integer)
            => I.Module -> Int -> I.Type
            -> IO [(I.Var, I.Block)]
sampleArray m len ty = repeatIO $ do
  (vars, blcks) <- unzip <$> replicateM len (head <$> sampleType m ty)
  let init = [ I.InitExpr ty (I.ExpVar v) | v <- vars ]
  (v, blck) <- mkLocal (I.TyArr len ty) (I.InitArray init True)
  return (v, concat blcks ++ blck)

repeatIO :: IO a -> IO [a]
repeatIO doThis = do
  x  <- doThis
  xs <- unsafeInterleaveIO (repeatIO doThis)
  return (x : xs)

lazyMapIO :: (a -> IO b) -> [a] -> IO [b]
lazyMapIO _ [] = return []
lazyMapIO f (a:as) = do
  b  <- f a
  bs <- unsafeInterleaveIO (lazyMapIO f as)
  return (b : bs)

zipLonger :: [[a]] -> [[b]] -> [([a], [b])]
zipLonger as bs = zip (as ++ repeat []) (bs ++ repeat [])