content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE NamedFieldPuns   #-}
  
  module Test.Generators where
  
  import Data.List (foldl')
  import Data.String (IsString (..))
  import Test.QuickCheck
  
  import Language.PureScript.Names (Ident (..), ModuleName (..), ProperName (..), Qualified (..), QualifiedBy (..))
  import Language.PureScript.PSString (PSString)
  import Language.PureScript.AST.SourcePos (SourceSpan (..), SourcePos (..))
  import Language.PureScript.AST (Literal (..))
  import Language.PureScript.CoreFn (Ann, Bind (..), Binder (..), CaseAlternative (..), Expr (..), Guard, ssAnn)
  
  import qualified Language.PureScript.DCE.Constants as C
  
  ann :: Ann
  ann = ssAnn (SourceSpan "src/Test.purs" (SourcePos 0 0) (SourcePos 0 0))
  
  genPSString :: Gen PSString
  genPSString = fromString <$> elements
    ["a", "b", "c", "d", "value0"]
  
  genProperName :: Gen (ProperName a)
  genProperName = ProperName <$> elements
    ["A", "B", "C", "D", "E"]
  
  genIdent :: Gen Ident
  genIdent = Ident <$> elements
    ["value0", "value1", "value2"]
  
  unusedIdents :: [Ident]
  unusedIdents =
    Ident <$> ["u1", "u2", "u3", "u4", "u5"]
  
  genUnusedIdent :: Gen Ident
  genUnusedIdent = elements unusedIdents
  
  genModuleName :: Gen ModuleName
  genModuleName = elements
    [ ModuleName "Data.Eq"
    , ModuleName "Data.Array"
    , ModuleName "Data.Maybe"
    , C.semigroup
    , C.unsafeCoerce
    , C.unit
    , C.semiring
    ]
  
  genQualifiedIdent :: Gen (Qualified Ident)
  genQualifiedIdent = oneof
    [ (Qualified . ByModuleName <$> genModuleName) <*> genIdent
    , return (Qualified (ByModuleName C.unit) (Ident "unit"))
    , return (Qualified (ByModuleName C.semiring) (Ident "add"))
    , return (Qualified (ByModuleName C.semiring) (Ident "semiringInt"))
    , return (Qualified (ByModuleName C.semiring) (Ident "semiringUnit"))
    , return (Qualified (ByModuleName C.maybeMod) (Ident "Just"))
    , return (Qualified (ByModuleName C.eqMod) (Ident "eq"))
    , return (Qualified (ByModuleName C.ring) (Ident "negate"))
    , return (Qualified (ByModuleName C.ring) (Ident "ringNumber"))
    , return (Qualified (ByModuleName C.ring) (Ident "unitRing"))
    ]
  
  genQualified :: Gen a -> Gen (Qualified a)
  genQualified gen = Qualified . ByModuleName <$> genModuleName <*> gen
  
  genLiteral :: Gen (Literal (Expr Ann))
  genLiteral = oneof
    [ NumericLiteral <$> arbitrary
    , StringLiteral  <$> genPSString
    , CharLiteral    <$> arbitrary
    , BooleanLiteral <$> arbitrary
    , ArrayLiteral . map unPSExpr <$> arbitrary
    , ObjectLiteral . map (\(k, v) -> (fromString k, unPSExpr v)) <$> arbitrary
    ]
  
  genLiteral' :: Gen (Expr Ann)
  genLiteral' = oneof
    [ Literal ann . NumericLiteral <$> arbitrary
    , Literal ann . StringLiteral <$> genPSString
    , Literal ann . BooleanLiteral <$> arbitrary
    , Literal ann . CharLiteral <$> arbitrary
    ]
  
  -- TODO: this generator is very frigile with size and at times can generate
  -- huge data.  We use size 4.  In particual it is very sensitive on the
  -- frequency of generating let expressions.
  genExpr :: Gen (Expr Ann)
  genExpr = sized go
    where
      go :: Int -> Gen (Expr Ann)
      go 0 = oneof
        [ Literal ann <$> genLiteral
        , Constructor ann <$> genProperName <*> genProperName <*> listOf genIdent
        , Var ann <$> genQualifiedIdent
        ]
      go n = frequency
        [ (3, Literal ann <$> genLiteral)
        , (3, Constructor ann <$> genProperName <*> genProperName <*> listOf genIdent)
        , (3, Var ann <$> genQualifiedIdent)
        , (4, Accessor ann <$> genPSString <*> scale succ genExpr)
        , (1, ObjectUpdate ann <$> genExpr <*> resize (max 3 (n - 1)) (listOf ((,) <$> genPSString <*> genExpr)))
        , (2, Abs ann <$> genIdent <*> scale succ genExpr)
        , (1, App ann <$> scale succ genExpr <*> scale succ genExpr)
        , (4, genApp)
        , (1, Case ann <$> resize (max 3 (n `div` 2)) (listOf genExpr) <*> resize (max 2 (n `div` 2)) (listOf (scale succ genCaseAlternative)))
        , (2, Let ann <$> listOf genBind <*> scale (`div` 2) genExpr)
        ]
  
  genCaseAlternative :: Gen (CaseAlternative Ann)
  genCaseAlternative = sized $ \n ->
    CaseAlternative <$> vectorOf n genBinder <*> genCaseAlternativeResult n
    where
    genCaseAlternativeResult :: Int -> Gen (Either [(Guard Ann, Expr Ann)] (Expr Ann))
    genCaseAlternativeResult n = oneof
      [ Left  <$> vectorOf n ((,) <$> resize n genExpr <*> resize n genExpr)
      , Right <$> resize n genExpr
      ]
  
  newtype PSBinder = PSBinder { unPSBinder :: Binder Ann }
    deriving Show
  
  genBinder :: Gen (Binder Ann)
  genBinder = sized go
    where
      go :: Int -> Gen (Binder Ann)
      go 0 = oneof
        [ return $ NullBinder ann
        , VarBinder ann <$> genIdent
        ]
      go _ = frequency
        [ (1, return $ NullBinder ann)
        , (2, LiteralBinder ann . ArrayLiteral  <$> listOf (scale succ genBinder))
        , (2, LiteralBinder ann . ObjectLiteral <$> listOf ((,) <$> genPSString <*> (scale succ genBinder)))
        , (3, ConstructorBinder ann <$> genQualified genProperName <*> genQualified genProperName <*> listOf (scale succ genBinder))
        , (3, NamedBinder ann <$> genIdent <*> scale succ genBinder)
        ]
  
  instance Arbitrary PSBinder where
    arbitrary =  PSBinder <$> resize 5 genBinder
  
    shrink (PSBinder (LiteralBinder _ (ArrayLiteral bs))) =
      (PSBinder . LiteralBinder ann . ArrayLiteral . map unPSBinder
      <$> (shrinkList shrink (PSBinder <$> bs)))
      ++ map PSBinder bs
    shrink (PSBinder (LiteralBinder _ (ObjectLiteral o))) =
      (PSBinder . LiteralBinder ann . ObjectLiteral
      <$> shrinkList (\(n, b) -> (n,) . unPSBinder <$> shrink (PSBinder b)) o)
      ++ map (PSBinder . snd) o
    shrink (PSBinder (ConstructorBinder _ tn cn bs)) =
      (PSBinder . ConstructorBinder ann tn cn . map unPSBinder
      <$> (shrinkList shrink (PSBinder <$> bs)))
      ++ map PSBinder bs
    shrink (PSBinder (NamedBinder _ n b)) =
      PSBinder b
      : (PSBinder . NamedBinder ann n . unPSBinder <$> shrink (PSBinder b))
    shrink _ = []
  
  prop_binderDistribution :: PSBinder -> Property
  prop_binderDistribution (PSBinder c) =
      classify True (show . depth $ c)
    $ tabulate "Binders" (cls c) True
    where
    cls NullBinder{}                 = ["NullBinder"]
    cls LiteralBinder{}              = ["LiteralBinder"]
    cls VarBinder{}                  = ["VarBinder"]
    cls (ConstructorBinder _ _ _ bs) = "ConstructorBinder" : concatMap cls bs
    cls (NamedBinder _ _ b)          = "NamedBinder" : cls b
  
    depth :: Binder a -> Int
    depth NullBinder{}                        = 1
    depth (LiteralBinder _ (ArrayLiteral bs)) = foldr (\b x -> depth b `max` x) 1 bs + 1
    depth (LiteralBinder _ (ObjectLiteral o)) = foldr (\(_, b) x -> depth b `max` x) 0 o + 1
    depth LiteralBinder{}                     = 1
    depth VarBinder{}                         = 1
    depth (ConstructorBinder _ _ _ bs)        = foldr (\b x -> depth b `max` x) 1 bs + 1
    depth (NamedBinder _ _ b)                 = depth b
  
  genBind :: Gen (Bind Ann)
  genBind = frequency
    [ (3, NonRec ann <$> gen  <*> (scale (`div` 2) genExpr))
    , (1, Rec <$> listOf ((\i e -> ((ann, i), e)) <$> gen <*> (scale (`div` 2) genExpr)))
    ]
    where
    gen = frequency [(3, genIdent), (2, genUnusedIdent)]
  
  newtype PSExpr a = PSExpr { unPSExpr :: Expr a }
    deriving Show
  
  -- Generate simple curried functions
  genApp :: Gen (Expr Ann)
  genApp =
    App ann
      <$> frequency
          [ (1, genApp)
          , (2, Var ann <$> genQualifiedIdent)
          ]
      <*> frequency
          [ (2, Var ann <$> genQualifiedIdent)
          , (3, genLiteral')
          ]
  
  instance Arbitrary (PSExpr Ann) where
    arbitrary = PSExpr <$> resize 4 genExpr
  
    shrink (PSExpr expr) = map PSExpr $ go expr
      where
      go :: Expr Ann -> [Expr Ann]
      go (Literal ann' (ArrayLiteral es)) =
        (Literal ann' . ArrayLiteral <$> shrinkList shrinkExpr es)
        ++ es
      go (Literal ann' (ObjectLiteral o)) =
        (Literal ann' . ObjectLiteral
        <$> shrinkList (\(n, e) -> (n,) <$> shrinkExpr e) o)
        ++ map snd o
      go (Accessor ann' n e) =
        e : (Accessor ann' n <$> shrinkExpr e)
      go (ObjectUpdate ann' e es) =
        e : map snd es
        ++ [ ObjectUpdate ann' e' es
           | e'  <- shrinkExpr e
           ]
        ++ [ ObjectUpdate ann' e es'
           | es' <- shrinkList (\(n, f) -> map (n,) $ shrinkExpr f) es
           ]
      go (Abs ann' n e) =
        let es = shrinkExpr e
        in e : es ++ map (Abs ann' n) es
      go (App ann' e f) =
        e : f : [ App ann' e' f  | e' <- shrinkExpr e ]
            ++  [ App ann' e  f' | f' <- shrinkExpr f ]
      go Var{} = []
      go (Case ann' es cs) =
           [ Case ann' es cs'
           | cs' <- shrinkList shrinkCaseAlternative cs
           ]
        ++ [ Case ann' es' cs
           | es' <- shrinkList shrinkExpr es
           ]
        ++ es
        ++ concatMap
            (\(CaseAlternative _ r) ->
              either
                (\es' -> map fst es' ++ map snd es')
                (\e' -> [e'])
                r
            )
            cs
      go (Let ann' bs e) =
        e : [ Let ann' bs  e' | e'  <- shrinkExpr e ]
         ++ [ Let ann' bs' e  | bs' <- shrinkList shrinkBind bs ]
      go _ = []
  
      shrinkCaseAlternative :: CaseAlternative Ann -> [CaseAlternative Ann]
      shrinkCaseAlternative (CaseAlternative bs r) =
           [ CaseAlternative bs r'
           | r' <-
               case r of
                 Right e  -> Right <$> shrinkExpr e
                 Left es' -> Left  <$> shrinkList
                                         (\(g, f) -> [(g, f') | f' <- shrinkExpr f]
                                                  ++ [(g', f) | g' <- shrinkExpr g]) es'
           ]
        ++ [ CaseAlternative bs' r
           | bs' <- shrinkList (\x -> [x]) bs
           ]
  
  shrinkExpr :: Expr Ann -> [Expr Ann]
  shrinkExpr = map unPSExpr . shrink . PSExpr
  
  shrinkBind :: Bind Ann -> [Bind Ann]
  shrinkBind (NonRec ann' n e) = NonRec ann' n <$> shrinkExpr e
  shrinkBind (Rec as) = Rec <$> shrinkList (\(x, e) -> map (x,) $ shrinkExpr e) as
  
  exprDepth :: Expr a -> Int
  exprDepth (Literal _ (ArrayLiteral es)) = foldl' (\x e -> exprDepth e `max` x) 1 es + 1
  exprDepth (Literal _ (ObjectLiteral o)) = foldl' (\x (_, e) -> exprDepth e `max` x) 1 o + 1
  exprDepth (Literal{})   = 1
  exprDepth Constructor{} = 1
  exprDepth (Accessor _ _ e) = 1 + exprDepth e
  exprDepth (ObjectUpdate _ e es) = 1 + exprDepth e + foldl' (\x (_, f) -> exprDepth f `max` x) 1 es
  exprDepth (Abs _ _ e) = 1 + exprDepth e
  exprDepth (App _ e f) = 1 + exprDepth e `max` exprDepth f
  exprDepth Var{}       = 1
  exprDepth (Case _ es cs) = 1 + foldl' (\x f -> exprDepth f `max` x) cdepth es
    where
      cdepth = foldl' (\x (CaseAlternative _ r) -> either (foldl' (\y (g, e) -> exprDepth g `max` exprDepth e `max` y) 1) exprDepth r `max` x) 1 cs
  exprDepth (Let _ bs e) = 1 + exprDepth e `max` foldl' (\x b -> binderExprDepth b `max` x) 0 bs
    where
      binderExprDepth :: Bind a -> Int
      binderExprDepth (NonRec _ _ e') = exprDepth e'
      binderExprDepth (Rec es') = foldl' (\x (_, e') -> x `max` exprDepth e') 0 es'
  
  prop_exprDistribution :: PSExpr Ann -> Property
  prop_exprDistribution (PSExpr e) =
      collect d
    $ tabulate "classify expressions" (cls e) True
    where
    d = exprDepth' e
    cls :: Expr a -> [String]
    cls Literal{}      = ["Literal"]
    cls Constructor{}  = ["Constructor"]
    cls Accessor{}     = ["Accessor"]
    cls ObjectUpdate{} = ["ObjectUpdate"]
    cls Abs{}          = ["Abs"]
    cls App{}          = ["App"]
    cls Var{}          = ["Var"]
    cls (Case _ _ cs)  = "Case" : foldl' (\x c -> clsCaseAlternative c ++ x) [] cs
      where
      clsCaseAlternative (CaseAlternative {caseAlternativeResult}) =
        either (foldl' (\x (g, f) -> cls g ++ cls f ++ x) []) cls caseAlternativeResult
    cls Let{}          = ["Let"]
  
    exprDepth' expr = case exprDepth expr of
      n | n < 10     -> n
        | n < 100   -> 10 * (n `div` 10)
        | n < 1000  -> 25 * (n `div` 25)
        | otherwise -> 100 * (n `div` 100)
  