content: {-# LANGUAGE CPP             #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  module Internal.TypeChecking.Generators where
  
  import Control.Applicative
  import Control.Monad.State
  import qualified Data.List as List (sort, nub)
  import Data.Text (Text)
  import qualified Data.Text as T
  import Data.Word
  
  import Agda.Syntax.Position
  import Agda.Syntax.Common
  import Agda.Syntax.Literal
  import Agda.Syntax.Fixity
  import Agda.Syntax.Internal as I
  import qualified Agda.Syntax.Concrete.Name as C
  
  import Agda.TypeChecking.Free
  import Agda.TypeChecking.Substitute
  
  import qualified Agda.Utils.VarSet as Set
  
  import Agda.Utils.Impossible
  
  import Internal.Helpers hiding ( Args )
  
  ------------------------------------------------------------------------------
  
  data TermConfiguration = TermConf
        { tcDefinedNames     :: [QName]
        , tcConstructorNames :: [QName]
        , tcProjectionNames  :: [QName]
        , tcFreeVariables    :: [Nat]
        , tcLiterals         :: UseLiterals
        , tcFrequencies      :: Frequencies
        , tcFixSize          :: Maybe Int
          -- ^ Maximum size of the generated element. When @Nothing@ this value
          --   is initialized from the 'Test.QuickCheck.size' parameter.
        , tcIsType           :: Bool
          -- ^ When this is true no lambdas, literals, or constructors are
          --   generated
        }
    deriving Show
  
  data Frequencies = Freqs
        { hiddenFreqs :: HiddenFreqs
        , elimFreqs   :: ElimFreqs
        , sortFreqs   :: SortFreqs
        , termFreqs   :: TermFreqs
        }
    deriving Show
  
  data TermFreqs = TermFreqs
        { varFreq :: Int
        , defFreq :: Int
        , conFreq :: Int
        , litFreq  :: Int
        , sortFreq :: Int
        , lamFreq  :: Int
        , piFreq   :: Int
        , funFreq  :: Int
        }
    deriving Show
  
  data ElimFreqs = ElimFreqs
        { applyFreq :: Int
        , projFreq  :: Int
        }
    deriving Show
  
  data HiddenFreqs = HiddenFreqs
        { hiddenFreq    :: Int
        , notHiddenFreq :: Int
        }
    deriving Show
  
  data SortFreqs = SortFreqs
        { setFreqs :: [Int]
        , propFreq :: Int
        }
    deriving Show
  
  defaultFrequencies :: Frequencies
  defaultFrequencies = Freqs
        { termFreqs   = TermFreqs   { varFreq = 24, defFreq = 8, conFreq = 8, litFreq = 1, sortFreq = 2, lamFreq = 10, piFreq = 5, funFreq = 5 }
        , elimFreqs   = ElimFreqs   { applyFreq = 9, projFreq = 1 }
        , hiddenFreqs = HiddenFreqs { hiddenFreq = 1, notHiddenFreq = 5 }
        , sortFreqs   = SortFreqs   { setFreqs = [3, 1], propFreq = 1 }
        }
  
  noProp :: TermConfiguration -> TermConfiguration
  noProp conf = conf { tcFrequencies = fq { sortFreqs = sfq { propFreq = 0 } } }
    where
      fq  = tcFrequencies conf
      sfq = sortFreqs fq
  
  data UseLiterals = UseLit
        { useLitInt    :: Bool
        , useLitFloat  :: Bool
        , useLitString :: Bool
        , useLitChar   :: Bool
        }
    deriving Show
  
  noLiterals :: UseLiterals
  noLiterals = UseLit False False False False
  
  fixSizeConf :: Int -> TermConfiguration -> TermConfiguration
  fixSizeConf n conf = conf { tcFixSize = Just n }
  
  resizeConf :: (Int -> Int) -> TermConfiguration -> TermConfiguration
  resizeConf f conf = conf { tcFixSize = fmap f $ tcFixSize conf}
  
  decrConf :: TermConfiguration -> TermConfiguration
  decrConf = resizeConf (flip (-) 1)
  
  divConf :: TermConfiguration -> Int -> TermConfiguration
  divConf conf k = resizeConf (`div` k) conf
  
  isTypeConf :: TermConfiguration -> TermConfiguration
  isTypeConf conf = conf { tcIsType = True }
  
  isntTypeConf :: TermConfiguration -> TermConfiguration
  isntTypeConf conf = conf { tcIsType = False }
  
  extendConf :: TermConfiguration -> TermConfiguration
  extendConf conf = conf { tcFreeVariables = 0 : map (1 +) (tcFreeVariables conf) }
  
  extendWithTelConf :: Telescope -> TermConfiguration -> TermConfiguration
  extendWithTelConf tel conf = foldr (const extendConf) conf (telToList tel)
  
  makeConfiguration :: [RawName] -> [RawName] -> [RawName] -> [Nat] -> TermConfiguration
  makeConfiguration ds cs ps vs = TermConf
    { tcDefinedNames     = defs
    , tcConstructorNames = cons
    , tcProjectionNames  = projs
    , tcFreeVariables    = List.sort $ List.nub vs
    , tcFrequencies      = defaultFrequencies
    , tcLiterals         = noLiterals
    , tcFixSize          = Nothing
    , tcIsType           = False
    }
    where
      (defs, cons, projs) = flip evalState 0 $ do
        (,,) <$> mapM mkName ds <*> mapM mkName cs <*> mapM mkName ps
  
      tick     = do x <- get; put (x + 1); return x
      mkName s = do
        n <- tick
        return $ QName { qnameModule = MName []
                       , qnameName   = Name
                          { nameId          = NameId n (ModuleNameHash 1)
                          , nameConcrete    = C.simpleName s
                          , nameCanonical   = C.simpleName s
                          , nameBindingSite = noRange
                          , nameFixity      = noFixity'
                          , nameIsRecordName = False
                          }
                        }
  
  class GenC a where
    genC :: TermConfiguration -> Gen a
  
  newtype YesType a   = YesType   { unYesType :: a     }
  newtype NoType  a   = NoType    { unNoType  :: a     }
  newtype VarName     = VarName   { unVarName :: Nat   }
  newtype DefName     = DefName   { unDefName :: QName }
  newtype ConName     = ConName   { unConName :: ConHead }
  newtype ProjName    = ProjName  { unProjName :: QName }
  newtype SizedList a = SizedList { unSizedList :: [a] }
  
  fixSize :: TermConfiguration -> Gen a -> Gen a
  fixSize conf g = sized $ \n -> resize (maybe n id $ tcFixSize conf) g
  
  instance GenC a => GenC (SizedList a) where
    genC conf = do
      n <- fixSize conf natural
      SizedList <$> vectorOf n (genC $ divConf conf n)
  
  instance GenC a => GenC [a] where
    genC conf = do
      n <- natural
      vectorOf n $ genC $ divConf conf n
  
  instance (GenC a, GenC b) => GenC (a, b) where
    genC conf = (,) <$> genC conf2 <*> genC conf2
      where
        conf2 = divConf conf 2
  
  instance GenC Range where
    genC _ = return noRange
  
  instance GenC Hiding where
    genC conf = frequency [ (hideF, return Hidden), (nohideF, return NotHidden) ]
      where
        HiddenFreqs {hiddenFreq = hideF, notHiddenFreq = nohideF } =
          hiddenFreqs $ tcFrequencies conf
  
  instance GenC a => GenC (Arg a) where
    genC conf = (\ (h, a) -> Arg (setHiding h defaultArgInfo) a) <$> genC conf
  
  instance GenC a => GenC (Dom a) where
    genC conf = (\ (h, a) -> defaultArgDom (setHiding h defaultArgInfo) a) <$> genC conf
  
  instance GenC a => GenC (Abs a) where
    genC conf = Abs x <$> genC (extendConf conf)
      where x = stringToArgName "x"
  
  instance GenC a => GenC (Elim' a) where
    genC conf = frequency [ (applyF, Apply <$> genC conf)
                          , (projF, Proj ProjSystem . unProjName <$> genC conf) ]
      where
        ElimFreqs {applyFreq = applyF, projFreq = projF } =
          elimFreqs $ tcFrequencies conf
  
  instance GenC DefName where
    genC conf = DefName  <$> do elements $ tcDefinedNames conf
  
  instance GenC ProjName where
    genC conf = ProjName <$> do elements $ tcProjectionNames conf
  
  genArgs :: TermConfiguration -> Gen Args
  genArgs conf = unSizedList <$> genC (isntTypeConf conf)
  
  genElims :: TermConfiguration -> Gen Elims
  genElims conf = unSizedList <$> genC (isntTypeConf conf)
  
  instance GenC Sort where
    genC conf = frequency $
      (propF, return (mkProp 0)) :
      zip setFs (map (return . mkType) [0..])
      where
        freq :: forall a. (Frequencies -> a) -> a
        freq f = f $ tcFrequencies conf
        setFs = freq (setFreqs . sortFreqs)
        propF = freq (propFreq . sortFreqs)
  
  instance GenC Char where
    genC _ = elements [' '..'~'] -- TODO
  
  instance GenC Double where
    genC _ = arbitrary
  
  instance GenC Integer where
    genC _ = arbitrary
  
  instance GenC Word64 where
    genC _ = arbitrary
  
  instance GenC Literal where
    genC conf = oneof $ concat $ zipWith gen useLits
                [ LitNat    <$> genC conf
                , LitWord64 <$> genC conf
                , LitFloat  <$> genC conf
                , LitString . T.pack <$> genC conf
                , LitChar   <$> genC conf
                ]
      where
        useLits = map ($ tcLiterals conf) [ useLitInt, useLitFloat, useLitString, useLitChar ]
  
        gen True  g = [g]
        gen False g = []
  
  instance GenC Telescope where
    genC conf = do
      n <- fixSize conf natural
      let confs = take n $ iterate extendConf (divConf conf n)
      telFromList <$> mapM genC confs
  
  instance GenC Type where
    genC conf = El <$> genC conf <*> genC (isTypeConf conf)
  
  instance GenC Term where
    genC conf = case tcFixSize conf of
        Nothing -> sized $ \n -> genC $ fixSizeConf n conf
        Just n | n <= 0    -> genLeaf
               | otherwise -> frequency
          [ (varF, genVar $ genElims conf)
          , (defF, genDef $ genElims conf)
          , (conF, genCon $ genArgs conf)
          , (litF,  Lit <$> genC conf)
          , (sortF, Sort <$> genC conf)
          , (lamF,  genLam)
          , (piF,   genPi)
          ]
      where
        defs    = tcDefinedNames conf
        cons    = tcConstructorNames conf
        vars    = tcFreeVariables conf
        freq f  = f $ tcFrequencies conf
        isType  = tcIsType conf
        useLits = map ($ tcLiterals conf) [ useLitInt, useLitFloat, useLitString, useLitChar ]
  
        varF  | null vars = 0
              | otherwise = freq (varFreq . termFreqs)
        defF  | null defs = 0
              | otherwise = freq (defFreq . termFreqs)
        conF  | null cons || isType = 0
              | otherwise           = freq (conFreq . termFreqs)
        litF  | or useLits && not isType = freq (litFreq . termFreqs)
              | otherwise             = 0
        lamF  | isType    = 0
              | otherwise = freq (lamFreq  . termFreqs)
        sortF = freq (sortFreq . termFreqs)
        piF   = freq (piFreq   . termFreqs)
  
        genLam :: Gen Term
        genLam = Lam <$> (flip setHiding defaultArgInfo <$> genC conf) <*> genC (isntTypeConf $ decrConf conf)
  
        genPi :: Gen Term
        genPi = uncurry Pi <$> genC conf
  
        genVar, genDef :: Gen Elims -> Gen Term
        genVar args = Var <$> elements vars <*> args
        genDef args = Def <$> elements defs <*> args
  
        genCon :: Gen Args -> Gen Term
        genCon args = Con <$> ((\ c -> ConHead c IsData Inductive []) <$> elements cons) <*> pure ConOSystem <*> (map Apply `fmap` args)
  
        genLeaf :: Gen Term
        genLeaf = frequency
          [ (varF, genVar $ return [])
          , (defF, genDef $ return [])
          , (conF, genCon $ return [])
          , (litF,  Lit  <$> genC conf)
          , (sortF, Sort <$> genC conf)
          ]
  
  -- | Only generates default configurations. Names and free variables varies.
  genConf :: Gen TermConfiguration
  genConf = do
    ds <- listOf $ elements defs
    cs <- listOf $ elements cons
    ps <- listOf1 $ elements projs
    vs <- listOf natural
    return $ makeConfiguration ds cs ps vs
    where
      defs = [ stringToRawName [c] | c <- ['a'..'n'] ++ ['r'..'z'] ]
      cons = [ stringToRawName [c] | c <- ['A'..'Z'] ]
      projs= [ stringToRawName [c] | c <- ['o'..'q'] ]
  
  instance Arbitrary TermConfiguration where
    arbitrary   = genConf
  
  -- Shrinking --------------------------------------------------------------
  
  class ShrinkC a where
    type ShrinksTo a
    shrinkC  :: TermConfiguration -> a -> [ShrinksTo a]
    noShrink :: a -> ShrinksTo a
  
  instance ShrinkC a => ShrinkC (YesType a) where
    type ShrinksTo (YesType a) = ShrinksTo a
    shrinkC conf (YesType x) = shrinkC (isTypeConf conf) x
    noShrink (YesType x) = noShrink x
  
  instance ShrinkC a => ShrinkC (NoType a) where
    type ShrinksTo (NoType a) = ShrinksTo a
    shrinkC conf (NoType x) = shrinkC (isntTypeConf conf) x
    noShrink (NoType x) = noShrink x
  
  instance ShrinkC a => ShrinkC [a] where
    type ShrinksTo [a] = [ShrinksTo a]
    noShrink        = map noShrink
    shrinkC conf xs = noShrink (removeChunks xs) ++ shrinkOne xs
     where
      -- Code stolen from Test.QuickCheck.Arbitrary
      removeChunks xs = rem (length xs) xs
       where
        rem 0 _  = []
        rem 1 _  = [[]]
        rem n xs = xs1
                 : xs2
                 : ( [ xs1' ++ xs2 | xs1' <- rem n1 xs1, not (null xs1') ]
               `ilv` [ xs1 ++ xs2' | xs2' <- rem n2 xs2, not (null xs2') ]
                   )
         where
          n1  = n `div` 2
          xs1 = take n1 xs
          n2  = n - n1
          xs2 = drop n1 xs
  
          []     `ilv` ys     = ys
          xs     `ilv` []     = xs
          (x:xs) `ilv` (y:ys) = x : y : (xs `ilv` ys)
  
      shrinkOne []     = []
      shrinkOne (x:xs) = [ x' : noShrink xs | x'  <- shrinkC conf x ]
                      ++ [ noShrink x : xs' | xs' <- shrinkOne xs ]
  
  instance (ShrinkC a, ShrinkC b) => ShrinkC (a, b) where
    type ShrinksTo (a, b) = (ShrinksTo a, ShrinksTo b)
    noShrink (x, y) = (noShrink x, noShrink y)
    shrinkC conf (x, y) =
      [ (x', noShrink y) | x' <- shrinkC conf x ] ++
      [ (noShrink x, y') | y' <- shrinkC conf y ]
  
  instance ShrinkC VarName where
    type ShrinksTo VarName = Nat
    shrinkC conf (VarName x) = [ y | y <- tcFreeVariables conf, y < x ]
    noShrink = unVarName
  
  instance ShrinkC DefName where
    type ShrinksTo DefName = QName
    shrinkC conf (DefName c) = takeWhile (/= c) $ tcDefinedNames conf
    noShrink = unDefName
  
  instance ShrinkC ConName where
    type ShrinksTo ConName = ConHead
    shrinkC conf (ConName (ConHead{conName = c})) = map (\ c -> ConHead c IsData Inductive []) $ takeWhile (/= c) $ tcConstructorNames conf
    noShrink = unConName
  
  instance ShrinkC Literal where
    type ShrinksTo Literal = Literal
    shrinkC _ (LitNat 0) = []
    shrinkC conf l         = LitNat 0 : case l of
        LitNat    n -> LitNat    <$> shrink n
        LitWord64 n -> LitWord64 <$> shrink n
        LitString s -> LitString . T.pack <$> shrinkC conf (T.unpack s)
        LitChar   c -> LitChar   <$> shrinkC conf c
        LitFloat  x -> LitFloat  <$> shrink x
        LitQName  x -> []
        LitMeta{}     -> []
    noShrink = id
  
  instance ShrinkC Char where
    type ShrinksTo Char = Char
    shrinkC _ 'a' = []
    shrinkC _ _   = ['a']
    noShrink = id
  
  instance ShrinkC Hiding where
    type ShrinksTo Hiding = Hiding
    shrinkC _ Hidden     = [NotHidden]
    shrinkC _ Instance{} = [NotHidden]
    shrinkC _ NotHidden  = []
    noShrink = id
  
  instance ShrinkC a => ShrinkC (Abs a) where
    type ShrinksTo (Abs a) = Abs (ShrinksTo a)
    shrinkC conf (NoAbs s x) = NoAbs s <$> shrinkC conf x
    shrinkC conf (Abs   s x) = Abs s <$> shrinkC (extendConf conf) x
    noShrink = fmap noShrink
  
  instance ShrinkC a => ShrinkC (Arg a) where
    type ShrinksTo (Arg a) = Arg (ShrinksTo a)
    shrinkC conf (Arg info x) = (\ (h,x) -> Arg (setHiding h info) x) <$> shrinkC conf (argInfoHiding info, x)
    noShrink = fmap noShrink
  
  instance ShrinkC a => ShrinkC (Dom a) where
    type ShrinksTo (Dom a) = Dom (ShrinksTo a)
    shrinkC conf dom@Dom{domInfo = info,unDom = x} = (\ (h,x) -> x <$ dom{domInfo = (setHiding h info)}) <$> shrinkC conf (argInfoHiding info, x)
    noShrink = fmap noShrink
  
  instance ShrinkC a => ShrinkC (Blocked a) where
    type ShrinksTo (Blocked a) = Blocked (ShrinksTo a)
    shrinkC conf (Blocked m x)    = Blocked m <$> shrinkC conf x
    shrinkC conf (NotBlocked r x) = NotBlocked r <$> shrinkC conf x
    noShrink = fmap noShrink
  
  instance ShrinkC a => ShrinkC (Elim' a) where
    type ShrinksTo (Elim' a) = Elim' (ShrinksTo a)
    shrinkC conf (Apply a) = Apply <$> shrinkC conf a
    shrinkC conf (IApply x y a) = IApply <$> shrinkC conf x <*> shrinkC conf y <*> shrinkC conf a
    shrinkC conf Proj{}    = []
    noShrink = fmap noShrink
  
  instance ShrinkC Sort where
    type ShrinksTo Sort = Sort
    shrinkC conf s = mkProp 0 : case s of
      Type n     -> [] -- No Level instance yet -- Type <$> shrinkC conf n
      SSet l     -> []
      Prop{}     -> []
      Inf f _    -> []
      SizeUniv   -> []
      LockUniv   -> []
      LevelUniv   -> []
      IntervalUniv -> []
      PiSort a s1 s2 -> __IMPOSSIBLE__
      FunSort s1 s2 -> __IMPOSSIBLE__
      UnivSort s -> __IMPOSSIBLE__
      MetaS x es -> __IMPOSSIBLE__
      DefS d es -> __IMPOSSIBLE__
      DummyS{} -> __IMPOSSIBLE__
    noShrink = id
  
  instance ShrinkC Telescope where
    type ShrinksTo Telescope = Telescope
    shrinkC conf EmptyTel          = []
    shrinkC conf (ExtendTel a tel) =
      killAbs tel : (uncurry ExtendTel <$> shrinkC conf (a, tel))
    noShrink = id
  
  instance ShrinkC Type where
    type ShrinksTo Type = Type
    shrinkC conf (El s t) = uncurry El <$> shrinkC conf (s, YesType t)
    noShrink = id
  
  instance ShrinkC Term where
    type ShrinksTo Term = Term
    shrinkC conf DontCare{}  = []
    shrinkC conf Dummy{}     = []
    shrinkC conf t           = filter validType $ case t of
      Var i es     -> map unArg (argsFromElims es) ++
                      (uncurry Var <$> shrinkC conf (VarName i, NoType es))
      Def d es     -> map unArg (argsFromElims es) ++
                      (uncurry Def <$> shrinkC conf (DefName d, NoType es))
      Con c ci es  -> map unArg (argsFromElims es) ++
                      ((\(c,vs) -> Con c ci vs) <$> shrinkC conf (ConName c, NoType es))
      Lit l        -> Lit <$> shrinkC conf l
      Level l      -> [] -- TODO
      Lam info b   -> killAbs b : ((\(h,x) -> Lam (setHiding h defaultArgInfo) x)
                                   <$> shrinkC conf (argInfoHiding info, b))
      Pi a b       -> unEl (unDom a) : unEl (killAbs b) :
                      (uncurry Pi <$> shrinkC conf (a, b))
      Sort s       -> Sort <$> shrinkC conf s
      MetaV m es   -> map unArg (argsFromElims es) ++
                      (MetaV m <$> shrinkC conf (NoType es))
  #if __GLASGOW_HASKELL__ < 900
      DontCare _   -> __IMPOSSIBLE__
      Dummy{}      -> __IMPOSSIBLE__
  #endif
      where
        validType t
          | not (tcIsType conf) = True
          | otherwise         = case t of
              Con{} -> False
              Lam{} -> False
              Lit{} -> False
              _     -> True
    noShrink = id
  
  killAbs :: KillVar a => Abs a -> a
  killAbs (Abs   _ x) = killVar 0 x
  killAbs (NoAbs _ x) = x
  
  class KillVar a where
    killVar :: Nat -> a -> a
  
  instance KillVar Term where
    killVar i t = case t of
      Var j args | j == i    -> DontCare (Var j [])
                 | j >  i    -> Var (j - 1) $ killVar i args
                 | otherwise -> Var j       $ killVar i args
      Def c args             -> Def c       $ killVar i args
      Con c ci args          -> Con c ci    $ killVar i args
      Lit l                  -> Lit l
      Level l                -> Level l -- TODO
      Sort s                 -> Sort s
      Lam h b                -> Lam h       $ killVar i b
      Pi a b                 -> uncurry Pi  $ killVar i (a, b)
      MetaV m args           -> MetaV m     $ killVar i args
      DontCare mv            -> DontCare    $ killVar i mv
      Dummy{}                -> t
  
  instance KillVar Type where
    killVar i (El s t) = El s $ killVar i t
  
  instance KillVar Telescope where
    killVar i EmptyTel          = EmptyTel
    killVar i (ExtendTel a tel) = uncurry ExtendTel $ killVar i (a, tel)
  
  instance KillVar a => KillVar (Elim' a) where
    killVar i = fmap (killVar i)
  
  instance KillVar a => KillVar (Arg a) where
    killVar i = fmap (killVar i)
  
  instance KillVar a => KillVar (Dom a) where
    killVar i = fmap (killVar i)
  
  instance KillVar a => KillVar (Abs a) where
    killVar i = fmap (killVar (i + 1))
  
  instance KillVar a => KillVar [a] where
    killVar i = map (killVar i)
  
  instance KillVar a => KillVar (Maybe a) where
    killVar i = fmap (killVar i)
  
  instance (KillVar a, KillVar b) => KillVar (a, b) where
    killVar i (x, y) = (killVar i x, killVar i y)
  
  -- Tests ------------------------------------------------------------------
  
  isWellScoped :: Free a => TermConfiguration -> a -> Bool
  isWellScoped conf t = allFreeVars t `Set.isSubsetOf` Set.fromList (tcFreeVariables conf)
  
  -- | Check that the generated terms don't have any out of scope variables.
  prop_wellScopedVars :: TermConfiguration -> Property
  prop_wellScopedVars conf =
    forAllShrink (genC conf) (shrinkC conf) $ \t ->
    isWellScoped conf (t :: Term)
  
  ------------------------------------------------------------------------
  -- * All tests
  ------------------------------------------------------------------------
  
  -- Template Haskell hack to make the following $allProperties work
  -- under ghc-7.8.
  return [] -- KEEP!
  
  -- | All tests as collected by 'allProperties'.
  --
  -- Using 'allProperties' is convenient and superior to the manual
  -- enumeration of tests, since the name of the property is added
  -- automatically.
  
  tests :: TestTree
  tests = testProperties "Internal.TypeChecking.Generators" $allProperties
  