content: {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  
  module ClosedWellTyped (
      ClosedWellTyped(..)
      ) where
  
  import Control.Applicative hiding (Const)
  import Control.Monad (guard)
  import Data.Text.Lazy (Text)
  import Data.Void (Void)
  import Control.Monad.State.Lazy (MonadState, StateT)
  import Control.Monad.Trans.Class (lift)
  import Morte.Core
  import Prelude hiding (pi)
  import Test.QuickCheck (Arbitrary, Gen)
  
  import qualified Control.Monad.State.Lazy as State
  import qualified Data.Text.Lazy           as Text
  import qualified Test.QuickCheck          as QuickCheck
  
  newtype ClosedWellTyped = ClosedWellTyped { unClosedWellTyped :: Expr Void }
      deriving Show
  
  data Env = Env
      { bindings :: [(Var, Expr Void)]
      , uniques  :: [Text]
      }
  
  newtype M a = M { unM :: StateT Env Gen a }
      deriving (Functor, Applicative, Monad, MonadState Env)
  
  liftGen :: Gen a -> M a
  liftGen m = M (lift m)
  
  evalM :: M a -> Env -> Gen a
  evalM m = State.evalStateT (unM m)
  
  instance Arbitrary ClosedWellTyped
    where
      arbitrary = QuickCheck.sized rndExpr
  
  rndExpr :: Int -> Gen ClosedWellTyped
  rndExpr n = fmap (ClosedWellTyped . fst) (evalM closed (initEnv n))
  
  initEnv :: Int -> Env
  initEnv n = Env
      { bindings = []
      , uniques  = map (Text.pack . show) [1..n]
      }
  
  extend :: Var -> Expr Void -> M ()
  extend x t = State.modify (\env -> env { bindings = (x, t) : bindings env })
  
  select :: [(Int, M (Expr Void, Expr Void), Env -> Bool)] -> M (Expr Void, Expr Void)
  select wgps = do
      env <- State.get
      m   <- liftGen (QuickCheck.frequency (do
          (weight, generator, predicate) <- wgps
          guard (predicate env)
          return (weight, return generator) ))
      m
  
  scope :: M a -> M a
  scope f = do
      env <- State.get
      liftGen (evalM f env)
  
  matching :: Eq b => b -> [(a, b)] -> Bool
  matching t = any ((t ==) . snd)
  
  moreThan :: Int -> [a] -> Bool
  moreThan n = not . null . drop n
  
  piFilter :: Expr Void -> Expr Void -> Bool
  piFilter t (Pi _ _A _) = _A == t
  piFilter _  _          = False
  
  closed :: M (Expr Void, Expr Void)
  closed =
      select
          [ (20, typeOrKind                       , \_ -> True          )
          , (50, lam (scope typeOrKind) termOrType, moreThan 0 . uniques)
          , (30, app                              , moreThan 1 . uniques)
          ]
  
  termOrType :: M (Expr Void, Expr Void)
  termOrType =
      select
          [ ( 5, type_                            , moreThan 0 . uniques )
          , (50, lam (scope typeOrKind) termOrType, moreThan 0 . uniques )
          , (25, var                              , moreThan 0 . bindings)
          , (20, app                              ,
              \e  ->  (null       (bindings e) && moreThan 1 (uniques e))
                  ||  (moreThan 0 (bindings e) && moreThan 0 (uniques e)) )
          ]
  
  typeOrKind :: M (Expr Void, Expr Void)
  typeOrKind =
      select
          [ (15, return (Const Star,Const Box)   , \_ -> True                      )
          , (20, varWith (== Const Star)         , matching (Const Star) . bindings)
          , (15, pi (scope typeOrKind) typeOrKind, moreThan 0            . uniques )
          ]
  
  varWith :: (Expr Void -> Bool) -> M (Expr Void, Expr Void)
  varWith f = do
      bEnv <- State.gets bindings
      liftGen (QuickCheck.elements (do
          (x, y) <- bEnv
          guard (f y)
          return (Var x, y) ))
  
  var :: M (Expr Void, Expr Void)
  var = varWith (\_ -> True)
  
  type_ :: M (Expr Void, Expr Void)
  type_ =
      select
          [ (20, varWith (== Const Star)     , matching (Const Star) . bindings)
          , (15, pi (scope typeOrKind) type_ , moreThan 0            . uniques )
          ]
  
  fresh :: M Text
  fresh = do
      env <- State.get
      let x:xs = uniques env
      State.put (env { uniques = xs })
      return x
  
  lam :: M (Expr Void, Expr Void)
      -> M (Expr Void, Expr Void)
      -> M (Expr Void, Expr Void)
  lam _AGen bGen = do
      x <- fresh
      (_A, _) <- _AGen
      extend (V x 0) _A
      (b, bType) <- bGen
      return (Lam x _A b, Pi x _A bType)
  
  pi  :: M (Expr Void, Expr Void)
      -> M (Expr Void, Expr Void)
      -> M (Expr Void, Expr Void)
  pi _AGen _BGen = do
      x <- fresh
      (_A, _) <- _AGen
      extend (V x 0) _A
      (_B, _BType) <- _BGen
      return (Pi x _A _B, _BType)
  
  app :: M (Expr Void, Expr Void)
  app = do
      (_N, _A       ) <- scope termOrType
      let genA = return (_A, Const Star)
      ~(f , Pi x _ _B) <- do
          scope
              (select
                  [ (40, lam genA termOrType  , moreThan 0              . uniques )
                  , (20, varWith (piFilter _A), any (piFilter _A . snd) . bindings)
                  ] )
      return (App f _N, subst x 0 _N _B)
  