{-
        



        Fun with Data Kinds

            Ulf Norell

           wg2.8 Aussois



-}

-- Not quite Haskell98
{-# LANGUAGE GADTs, DataKinds, PolyKinds, KindSignatures, TypeOperators,
             StandaloneDeriving, MultiParamTypeClasses,
             RankNTypes, UndecidableInstances, TypeFamilies,
             ScopedTypeVariables, MagicHash, TypeHoles, ImplicitParams
  #-}

module FunWithDataKinds where

import Control.Arrow ((***), (&&&))
import Control.Applicative
import Control.Monad
import Control.Monad.State
import Data.Monoid hiding ((<>))
import Test.QuickCheck
import Text.PrettyPrint as P
import Data.Proxy
import GHC.TypeLits
import Data.Void
import Data.Function
import GHC.Prim
import Debug.Trace
import Data.Char
import System.Random

-- As well all know state monads are not allowed to
-- increase the state.

prop_monad :: Tm (State Int ()) -> Bool
prop_monad (_ :~ m) = execState m 100 <= 100

-- *FunWithDataKinds> quickCheck prop_monad
-- *** Failed! Falsifiable (after 64 tests and 1 shrink):     
-- get >>= (\a -> put (a + a))

-- Or maybe that's not right...

-- Modelling type correct terms

data CXT a = NIL | CXT a ::> a

infixl 5 :>, ::>
data Cxt :: (k -> *) -> CXT k -> * where
  Nil  :: Cxt ty NIL
  (:>) :: Cxt ty ts -> ty t -> Cxt ty (ts ::> t)

data Var :: CXT k -> k -> * where
  Vz :: Var (g ::> a) a
  Vs :: Var g a -> Var (g ::> b) a

(!) :: Cxt ty cxt -> Var cxt a -> ty a
(cxt :> a) ! Vz   = a
(cxt :> a) ! Vs i = cxt ! i

data TYPE   = TCON Symbol | TCON1 Symbol TYPE   | TYPE   ::-> TYPE

-- Constants can be polymorphic (and thus have schemes). We fix two type
-- variables X and Y (for reasons of lack of full dependent types).
data SCHEME = SCON Symbol | SCON1 Symbol SCHEME | SCHEME ::=> SCHEME
            | X | Y

-- Symbol is the kind of type level strings.
data Name :: Symbol -> * where
  Name :: KnownSymbol c => Proxy c -> Name c

data Type :: TYPE -> * where
  TCon   :: Name c -> Type (TCON c)
  TCon1  :: Name c -> Type a -> Type (TCON1 c a)
  (:->)  :: Type a -> Type b -> Type (a ::-> b)

data Scheme :: SCHEME -> * where
  SVarX :: Scheme X
  SVarY :: Scheme Y
  SCon  :: Name c -> Scheme (SCON c)
  SCon1 :: Name c -> Scheme a -> Scheme (SCON1 c a)
  (:=>) :: Scheme a -> Scheme b -> Scheme (a ::=> b)

infixr 7 ::->, :->, :=>, ::=>

-- Singelton type class

class Singleton (s :: k -> *) (x :: k) where
  it :: s x

class WithSingleton s where
  withSingleton :: s x -> (Singleton s x => a) -> a

instance Singleton  (Cxt f) NIL where
  it = Nil
instance (Singleton (Cxt f) cxt, Singleton f a) =>
         Singleton (Cxt f) (cxt ::> a) where
  it = it :> it

instance WithSingleton ty => WithSingleton (Cxt ty) where
  withSingleton Nil k = k
  withSingleton (xs :> x) k = withSingleton xs $ withSingleton x k

instance KnownSymbol c => Singleton Name c where it = Name Proxy
instance Singleton Name c => Singleton Type (TCON c) where it = TCon it
instance (Singleton Name c, Singleton Type a) => Singleton Type (TCON1 c a) where it = TCon1 it it
instance (Singleton Type a, Singleton Type b) => Singleton Type (a ::-> b) where it = it :-> it

instance WithSingleton Name where
  withSingleton Name{} k = k

instance WithSingleton Type where
  withSingleton a k =
    case a of
      TCon c    -> withSingleton c k
      TCon1 c a -> withSingleton c $ withSingleton a k
      a :-> b   -> withSingleton a $ withSingleton b k

instance Singleton Name c => Singleton Scheme (SCON c) where it = SCon it
instance (Singleton Name c, Singleton Scheme a) => Singleton Scheme (SCON1 c a) where it = SCon1 it it
instance (Singleton Scheme a, Singleton Scheme b) => Singleton Scheme (a ::=> b) where it = it :=> it
instance Singleton Scheme X where it = SVarX
instance Singleton Scheme Y where it = SVarY

instance WithSingleton Scheme where
  withSingleton a k =
    case a of
      SVarX     -> k
      SVarY     -> k
      SCon c    -> withSingleton c k
      SCon1 c a -> withSingleton c $ withSingleton a k
      a :=> b   -> withSingleton a $ withSingleton b k

-- Terms

data Term :: CXT SCHEME -> TYPE -> * where
  Var  :: Var cxt a -> Type x -> Type y -> Term cxt (Inst a x y)
  (:@) :: Term cxt (a ::-> b) -> Term cxt a -> Term cxt b
  Lam  :: Type a -> Term (cxt ::> Mono a) b -> Term cxt (a ::-> b)

-- Some useful functions on schemes and types

-- Closed type families
type family Inst  (s :: SCHEME) (x :: TYPE)   (y :: TYPE)   :: TYPE where
  Inst  X           x y = x
  Inst  Y           x y = y
  Inst  (SCON c)    x y = TCON c
  Inst  (SCON1 c a) x y = TCON1 c (Inst a x y)
  Inst  (a ::=> b)  x y = Inst a x y ::-> Inst b x y

type family InstS (s :: SCHEME) (x :: SCHEME) (y :: SCHEME) :: SCHEME where
  InstS X           x y = x
  InstS Y           x y = y
  InstS (SCON c)    x y = SCON c
  InstS (SCON1 c a) x y = SCON1 c (InstS a x y)
  InstS (a ::=> b)  x y = InstS a x y ::=> InstS b x y

type family Mono  (a :: TYPE) :: SCHEME where
  Mono  (TCON c)    = SCON c
  Mono  (TCON1 c a) = SCON1 c (Mono a)
  Mono  (a ::-> b)  = Mono a ::=> Mono b

-- Corresponding value level functions

inst  :: Scheme s -> Type x   -> Type y   -> Type   (Inst s x y)
instS :: Scheme s -> Scheme x -> Scheme y -> Scheme (InstS s x y)
mono  :: Type a -> Scheme (Mono a)

mono a = case a of
  TCon c    -> SCon c
  TCon1 c a -> SCon1 c (mono a)
  a :-> b   -> mono a :=> mono b

inst s x y = case s of
  SVarX     -> x
  SVarY     -> y
  SCon c    -> TCon c
  SCon1 c s -> TCon1 c (inst s x y)
  a :=> b   -> inst a x y :-> inst b x y

instS s x y = case s of
  SVarX     -> x
  SVarY     -> y
  SCon c    -> SCon c
  SCon1 c s -> SCon1 c (instS s x y)
  a :=> b   -> instS a x y :=> instS b x y

-- Type equality proofs

data (:=:) :: k -> k -> * where
  Refl :: a :=: a
  -- really: a ~ b => a :=: b

rewrite :: a :=: b -> (a ~ b => c) -> c
rewrite Refl x = x

infixr 0 ~$
(~$) = rewrite

class Equal (f :: k -> *) where
  (=?=) :: f a -> f b -> [a :=: b]
        -- MonadZero m => f a -> f b -> m (a :=: b)

instance Equal Name where
  Name x =?= Name y
    | symbolVal x == symbolVal y = [ unsafeCoerce# Refl ]
    | otherwise                  = []

instance Equal Type where
  TCon c    =?= TCon c'    = [ Refl | Refl <- c =?= c' ]
  TCon1 c a =?= TCon1 c' b = [ Refl | Refl <- c =?= c', Refl <- a =?= b ]
  (a :-> b) =?= (c :-> d)  = [ Refl | Refl <- a =?= c,  Refl <- b =?= d ]
  _         =?= _          = []

-- We're going to need some lemmas

lemInstSId   :: Scheme s -> InstS s X Y :=: s

lemInstMono  :: Type a -> Type x -> Type y ->
                Inst (Mono a) x y :=: a

lemInstSMono :: Type a -> Scheme x -> Scheme y ->
                InstS (Mono a) x y :=: Mono a

lemInstS     :: Scheme s -> Scheme x -> Scheme y -> Scheme z -> Scheme w ->
                InstS (InstS s x y) z w :=:
                InstS s (InstS x z w) (InstS y z w)

lemInst      :: Scheme s -> Scheme x -> Scheme y -> Type z -> Type w ->
                Inst (InstS s x y) z w :=:
                Inst s (Inst x z w) (Inst y z w)

-- proofs
lemInstSMono TCon{}      _ _ = Refl
lemInstSMono (TCon1 _ a) x y = lemInstSMono a x y ~$ Refl
lemInstSMono (a :-> b)   x y = lemInstSMono a x y ~$
                               lemInstSMono b x y ~$ Refl

lemInstMono TCon{}      _ _ = Refl
lemInstMono (TCon1 _ a) x y = lemInstMono a x y ~$ Refl
lemInstMono (a :-> b)   x y = lemInstMono a x y ~$
                              lemInstMono b x y ~$ Refl

lemInstS s sx sy x y =
  case s of
    SVarX     -> Refl
    SVarY     -> Refl
    SCon c    -> Refl
    SCon1 c a -> lemInstS a sx sy x y ~$ Refl
    a :=> b   -> lemInstS a sx sy x y ~$
                 lemInstS b sx sy x y ~$ Refl

lemInst s sx sy x y =
  case s of
    SVarX     -> Refl
    SVarY     -> Refl
    SCon c    -> Refl
    SCon1 c a -> lemInst a sx sy x y ~$ Refl
    a :=> b   -> lemInst a sx sy x y ~$
                 lemInst b sx sy x y ~$ Refl

lemInstSId s =
  case s of
    SVarX     -> Refl
    SVarY     -> Refl
    SCon c    -> Refl
    SCon1 c a -> lemInstSId a ~$ Refl
    a :=> b   -> lemInstSId a ~$ lemInstSId b ~$ Refl

-- Semantics of types and schemes

--   Parameterized by a signature assigning types to type names.

type family TVal  sig (name :: Symbol) :: *
type family TVal1 sig (name :: Symbol) :: * -> *

-- TVal instances that hold for any signature
type instance TVal  sig "Int"  = Int
type instance TVal  sig "Bool" = Bool
type instance TVal  sig "()"   = ()
type instance TVal1 sig "[]"   = []

type family   El sig (a :: TYPE) :: *
type instance El sig (TCON c)    = TVal sig c
type instance El sig (TCON1 c a) = TVal1 sig c (El sig a)
type instance El sig (a ::-> b)  = El sig a -> El sig b

newtype ElS sig s =
  ElS { instantiate :: forall x y. Type x -> Type y ->
                                   El sig (Inst s x y) }

-- Interpreter for terms

eval :: forall sig cxt a. Cxt (ElS sig) cxt -> Term cxt a -> El sig a
eval env v =
  case v of
    Var x a b -> instantiate (env ! x) a b
    f :@ v    -> eval env f (eval env v)
    Lam a v   -> \x -> eval (env :> liftMono a x) v

liftMono :: Type a -> El sig a -> ElS sig (Mono a)
liftMono a v = ElS $ \x y -> lemInstMono a x y ~$ v

-- Some additional types that will come in handy

data Some :: (k -> *) -> * where
  Like :: f k -> Some f

data Spine :: CXT SCHEME -> TYPE -> * where
  Head :: Var g a -> Type x -> Type y -> Spine g (Inst a x y)
  Arg  :: Spine g (a ::-> b) -> Type a -> Spine g b

data SpineS :: CXT SCHEME -> SCHEME -> * where
  HeadS :: Var g a -> Scheme x -> Scheme y -> SpineS g (InstS a x y)
  ArgS  :: SpineS g (a ::=> b) -> Scheme a -> SpineS g b

-- Generating random terms

data TyConGen = TyConGen { genTyCon, genTyCon1 :: Gen (Some Name) }

genTerm :: forall g a. TyConGen -> Cxt Scheme g -> Type a -> Int ->
                       Gen (Term g a)
genTerm gc g a n = oneof $

    [ Lam b <$> genTerm gc (g :> mono b) c n
      | b :-> c <- [a] ] ++

    [ genApply a s | s <- matchingSpines g a,
                     n > 0 || simpleSpine a s ]

  where
    genApply :: Type a -> SpineS g (Mono a) -> Gen (Term g a)
    genApply a s = do
      Like x <- genType gc 2  -- Pick small types to instantiate
      Like y <- genType gc 2  -- unconstrained type vars
      lemInstMono a x y ~$
        genArgs n' (instSpine g s x y)
      where
        n' = argSize n (spineLen s)

    genArgs :: Int -> Spine g b -> Gen (Term g b)
    genArgs _ (Head x a b) = return (Var x a b)
    genArgs m (Arg s a)    = (:@) <$> genArgs m s <*> genTerm gc g a m

    argSize n 0 = 0
    argSize n 1 = max 0 (n - 1)
    argSize n k = div n k

genType :: TyConGen -> Int -> Gen (Some Type)
genType gc 0 = do
  Like c <- genTyCon gc
  return $ Like (TCon c)
genType gc n = oneof
  [ do Like c <- genTyCon gc
       return $ Like (TCon c)
  , do Like c <- genTyCon1 gc
       Like a <- genType gc (n - 1)
       return $ Like (TCon1 c a)
  , do Like a <- genType gc (div n 3)
       Like b <- genType gc (div (2 * n) 3)
       return $ Like (a :-> b)
  ]

genScheme :: TyConGen -> Int -> Gen (Some Scheme)
genScheme gc 0 =
  oneof [ pure $ Like SVarX
        , pure $ Like SVarY
        , do Like c <- genTyCon gc; return $ Like (SCon c)
        ]
genScheme gc n = oneof
  [ genScheme gc 0
  , do Like c <- genTyCon1 gc
       Like a <- genScheme gc (n - 1)
       return $ Like $ SCon1 c a
  , do Like a <- genScheme gc (div n 4)
       Like b <- genScheme gc (div (3 * n) 4)
       return $ Like $ a :=> b
  ]

simpleSpine :: Type a -> SpineS g b -> Bool
simpleSpine _ HeadS{}    = True
simpleSpine a (ArgS s b) = sizeS b < size a && simpleSpine a s
  where
    sizeS :: Scheme a -> Int
    sizeS SVarX       = 0
    sizeS SVarY       = 0
    sizeS SCon{}      = 0
    sizeS (SCon1 _ a) = 1 + sizeS a
    sizeS (a :=> b)   = 1 + sizeS a + sizeS b

    size :: Type a -> Int
    size TCon{}      = 0
    size (TCon1 _ a) = 1 + size a
    size (a :-> b)   = 1 + size a + size b

-- Compute spines for all variables in the context that can
-- construct something of the given type.

matchingSpines :: Cxt Scheme g -> Type a -> [SpineS g (Mono a)]
matchingSpines Nil      a = []
matchingSpines (g :> s) a =
    matches (g :> s) s a (lemInstSId s ~$ HeadS Vz SVarX SVarY) ++
    [ wkSpineS s | s <- matchingSpines g a ]
  where
    matches :: Cxt Scheme g -> Scheme s -> Type a -> SpineS g s -> [SpineS g (Mono a)]
    matches cxt s a sp =
      [ lemInstSMono a x y ~$ instSpineS cxt sp x y
      | Unified x y Refl <- unify s (mono a) ] ++
      case s of
        b :=> c -> matches cxt c a (sp `ArgS` b)
        _       -> []

-- Unification

data Unified :: SCHEME -> SCHEME -> * where
  Unified :: Scheme x -> Scheme y ->
             InstS a x y :=: InstS b x y -> Unified a b

unify :: Scheme a -> Scheme b -> [Unified a b]
unify SVarX SVarX        = [uTrivial]
unify SVarY SVarY        = [uTrivial]
unify SVarX a = do
  NoOccurs eq <- occursCheck SVarX a
  return $ lemSubstX a a
        ~$ eq SVarX a
        ~$ Unified a SVarY Refl
unify a SVarX = do
  NoOccurs eq <- occursCheck SVarX a
  return $ lemSubstX a a
        ~$ eq SVarX a
        ~$ Unified a SVarY Refl
unify a SVarY = do
  NoOccurs eq <- occursCheck SVarY a
  return $ lemSubstY a a
        ~$ eq SVarY a
        ~$ Unified SVarX a Refl
unify SVarY a = do
  NoOccurs eq <- occursCheck SVarY a
  return $ lemSubstY a a
        ~$ eq SVarY a
        ~$ Unified SVarX a Refl
unify (SCon c) (SCon c') = [uTrivial | Refl <- c =?= c']
unify (SCon1 c a) (SCon1 c' b) = do
  Refl           <- c =?= c'
  Unified x y eq <- unify a b
  eq ~$ return $ Unified x y Refl
unify (a1 :=> b1) (a2 :=> b2) = do
  Unified x1 y1 eq1 <- unify a1 a2
  Unified x2 y2 eq2 <- unify (instS b1 x1 y1) (instS b2 x1 y1)
  eq1 ~$ eq2
      ~$ lemInstS a1 x1 y1 x2 y2
      ~$ lemInstS a2 x1 y1 x2 y2
      ~$ lemInstS b1 x1 y1 x2 y2
      ~$ lemInstS b2 x1 y1 x2 y2
      ~$ return $ Unified (instS x1 x2 y2) (instS y1 x2 y2) Refl
unify _ _ = []

uTrivial :: Unified a a
uTrivial = Unified SVarX SVarY Refl

-- Occurs checking

-- z ∈ {X, Y}
data NoOccurs :: SCHEME -> SCHEME -> * where
  NoOccurs :: (forall x. Scheme z -> Scheme x -> Subst z x s :=: s) ->
              NoOccurs z s

type family Subst (z :: SCHEME) (s :: SCHEME) (t :: SCHEME) :: SCHEME where
  Subst X s X            = s
  Subst Y s Y            = s
  Subst X s Y            = Y 
  Subst Y s X            = X
  Subst z s (SCON c)     = SCON c
  Subst z s (SCON1 c t)  = SCON1 c (Subst z s t)
  Subst z s (t1 ::=> t2) = Subst z s t1 ::=> Subst z s t2

occursCheck :: Scheme z -> Scheme s -> [NoOccurs z s]
occursCheck SVarX SVarX  = []
occursCheck SVarX SVarY  = [NoOccurs $ \_ _ -> Refl]
occursCheck SVarY SVarX  = [NoOccurs $ \_ _ -> Refl]
occursCheck SVarY SVarY  = []
occursCheck z SCon{} = [NoOccurs $ \_ _ -> Refl]
occursCheck z (SCon1 c a) = do
  NoOccurs prf <- occursCheck z a
  return $ NoOccurs $ \_ x -> prf z x ~$ Refl
occursCheck z (a :=> b) = do
  NoOccurs prfa <- occursCheck z a
  NoOccurs prfb <- occursCheck z b
  return $ NoOccurs $ \_ x -> prfa z x ~$ prfb z x ~$ Refl

lemSubstX :: Scheme x -> Scheme s -> Subst X x s :=: InstS s x Y
lemSubstY :: Scheme y -> Scheme s -> Subst Y y s :=: InstS s X y

lemSubstX x s =
  case s of 
    SVarX     -> Refl
    SVarY     -> Refl
    SCon{}    -> Refl
    SCon1 _ a -> lemSubstX x a ~$ Refl
    a :=> b   -> lemSubstX x a ~$ lemSubstX x b ~$ Refl

lemSubstY y s =
  case s of 
    SVarX     -> Refl
    SVarY     -> Refl
    SCon{}    -> Refl
    SCon1 _ a -> lemSubstY y a ~$ Refl
    a :=> b   -> lemSubstY y a ~$ lemSubstY y b ~$ Refl

-- A signature defines a bunch of types and constants

class (IsCxt (SigCxt sig), Singleton (Cxt Scheme) (SigCxt sig)) =>
      Signature (sig :: *) where

  tyCons  :: Proxy sig -> [Some Name]
  tyCons1 :: Proxy sig -> [Some Name]

  type SigCxt sig :: CXT SCHEME
  sigEnv :: Proxy sig -> Cxt (Constant sig) (SigCxt sig)

data Constant tenv t =
  (:=) { cName :: String, cVal :: ElS tenv t }

tyConGen :: Signature sig => Proxy sig -> TyConGen
tyConGen p = TyConGen (elements $ tyCons p) (elements $ tyCons1 p)

-- Trickery. Don't ask.
(=:) :: String ->
        (forall x y. (?ty_x :: Type x, ?ty_y :: Type y) =>
                     El sig (Inst a x y)) ->
        Constant sig a
x =: v = x := ElS (\ty_x ty_y -> let ?ty_x = ty_x; ?ty_y = ty_y in v)

-- Closed terms. Just to get access to names of constants from
-- the signature when printing.

data Closed :: * -> TYPE -> * where
  Closed :: Signature sig => Term (SigCxt sig) a -> Closed sig a

-- evaluation and generation lifts without problems
evalClosed :: forall sig a. Signature sig => Closed sig a -> El sig a
evalClosed (Closed v) = eval env v
  where
    env :: Cxt (ElS sig) (SigCxt sig)
    env = mapCxt cVal $ sigEnv (Proxy :: Proxy sig)

mapCxt :: (forall a. f a -> g a) -> Cxt f cxt -> Cxt g cxt
mapCxt f Nil        = Nil
mapCxt f (cxt :> x) = mapCxt f cxt :> f x

genClosed :: forall sig a. Signature sig => Type a -> Int -> Gen (Closed sig a)
genClosed a n = Closed <$> genTerm tcg cxt a n
  where
    tcg = tyConGen (Proxy :: Proxy sig)
    cxt = it :: Cxt Scheme (SigCxt sig)

instance (Singleton Type a, Signature sig) => Arbitrary (Closed sig a) where
  arbitrary = sized $ genClosed it
  shrink = shrinkClosed it

-- A closed term and its value
data Glued :: * -> TYPE -> * where
  (:~) :: Signature sig => Closed sig a -> El sig a -> Glued sig a

instance (Singleton Type a, Signature sig) => Arbitrary (Glued sig a) where
  arbitrary = glued <$> arbitrary
  shrink (v :~ _) = [ glued v | v <- shrink v ]

glued :: Closed sig a -> Glued sig a
glued v@Closed{} = v :~ evalClosed v

-- Example signature --

data MonadSig (m :: * -> *)
data StateM a

type instance ToScheme (StateM a) = SCON1 "m" (ToScheme a)
type instance TVal1 (MonadSig m) "m" = m

instance MonadState Int m => Signature (MonadSig m) where
  type SigCxt (MonadSig m) =
    ToCxt (NIL ::> (StateM A -> (A -> StateM B) -> StateM B)
               ::> (A -> StateM A)
               ::> StateM Int ::> (Int -> StateM ())
               ::> (Int -> Int -> Int) ::> Int ::> ())
  tyCons  _ = [Like (it :: Name "Int"), Like (it :: Name "()")]
  tyCons1 _ = [Like (it :: Name "m")]
  sigEnv _ = Nil :> ">>="    =: (>>=)
                 :> "return" =: return
                 :> "get"    =: get
                 :> "put"    =: put
                 :> "+"      =: (+)
                 :> "1"      =: 1
                 :> "()"     =: ()

type Tm a = Glued (MonadSig (State Int)) (ToType a)

sampleGen :: Gen (Closed (MonadSig (State Int)) (ToType (Int -> StateM Int)))
sampleGen = arbitrary

{-

    End of the actual talk.

     Below: Bonus Material

-}

-- Conventient functions to convert from types in * to TYPE or SCHEME

type family ToType   a :: TYPE
type family ToScheme a :: SCHEME
type family ToCxt (xs :: CXT *) :: CXT SCHEME

type instance ToType Int        = TCON "Int"
type instance ToType Bool       = TCON "Bool"
type instance ToType [a]        = TCON1 "[]" (ToType a)
type instance ToType (a -> b)   = ToType a ::-> ToType b
type instance ToType (StateM a) = TCON1 "m" (ToType a)
type instance ToType ()         = TCON "()"
type instance ToType (State Int a) = TCON1 "m" (ToType a)

data A
data B

type instance ToScheme Int        = SCON "Int"
type instance ToScheme Bool       = SCON "Bool"
type instance ToScheme [a]        = SCON1 "[]" (ToScheme a)
type instance ToScheme (a -> b)   = ToScheme a ::=> ToScheme b
type instance ToScheme ()         = SCON "()"
type instance ToScheme A          = X
type instance ToScheme B          = Y

type instance ToCxt NIL = NIL
type instance ToCxt (as ::> a) = ToCxt as ::> ToScheme a

-- Shrinking --

shrinkVar :: Var g a -> [Some (Var g)]
shrinkVar Vz     = []
shrinkVar (Vs i) = case wkVar i of Like i -> Like i : shrinkVar i

shrinkSomeTerm :: IsCxt g => Cxt Scheme g -> Some (Term g) -> [Some (Term g)]
shrinkSomeTerm cxt (Like v) =
  case v of
    Var x a b -> [ Like (Var x a b) | Like x <- shrinkVar x ]
    f :@ v    -> Like f : Like v :
                 [ Like (f :@ v) | f <- shrinkTerm cxt (typeOf cxt f) f ] ++
                 [ Like (f :@ v) | v <- shrinkTerm cxt (typeOf cxt v) v ]
    Lam a v   -> [ Like v | v <- strengthen cxt (mono a) v ] ++
                 [ Like f | f :@ u <- [v], f <- strengthen cxt (mono a) f ] ++
                            -- eta contraction
                 [ Like (Lam a v) | Like v <- shrinkSomeTerm (cxt :> mono a) (Like v) ]

shrinkTerm :: IsCxt g => Cxt Scheme g -> Type a -> Term g a -> [Term g a]
shrinkTerm cxt a v = do
  Like u <- shrinkSomeTerm cxt (Like v)
  case typeOf cxt u =?= a of
    [Refl] -> [u]
    []     -> case u of
      Var x _ _ -> do
        -- If shrinking to a ill-typed variable, try to reunify the type
        -- arguments.
        Unified b c eq <- unify (cxt ! x) (mono a)
        Like z <- [baseType a]
        return $ eq
              ~$ lemInst (cxt ! x) b c z z
              ~$ lemInstMono a z z
              ~$ lemInstSMono a b c
              ~$ Var x (inst b z z) (inst c z z)
      _ -> []

wkVar :: Var g a -> Some (Var (g ::> b))
wkVar Vz     = Like Vz
wkVar (Vs i) = case wkVar i of Like j -> Like (Vs j)

infixl 5 ++, ++>

type family (++) (xs :: CXT f) (ys :: CXT f) :: CXT f
type instance xs ++ NIL = xs
type instance xs ++ (ys ::> y) = (xs ++ ys) ::> y

(++>) :: Cxt f xs -> Cxt f ys -> Cxt f (xs ++ ys)
xs ++> Nil       = xs
xs ++> (ys :> y) = (xs ++> ys) :> y

strengthenVar :: Cxt ty g1 -> Cxt ty g2 -> ty a ->
                 Var (g1 ::> a ++ g2) b -> [Var (g1 ++ g2) b]
strengthenVar g1 Nil       a Vz     = []
strengthenVar g1 (g2 :> b) a Vz     = [Vz]
strengthenVar g1 Nil       a (Vs x) = [x]
strengthenVar g1 (g2 :> b) a (Vs x) = [Vs x | x <- strengthenVar g1 g2 a x]

strength :: Cxt Scheme g1 -> Cxt Scheme g2 -> Scheme a ->
            Term (g1 ::> a ++ g2) b -> [Term (g1 ++ g2) b]
strength g1 g2 a v =
  case v of
    Var v x y -> [ Var v x y | v <- strengthenVar g1 g2 a v ]
    f :@ v    -> [ f :@ v    | f <- strength g1 g2 a f, v <- strength g1 g2 a v ]
    Lam b v   -> [ Lam b v   | v <- strength g1 (g2 :> mono b) a v ]

strengthen :: Cxt Scheme g -> Scheme a -> Term (g ::> a) b -> [Term g b]
strengthen cxt a v = strength cxt Nil a v

baseType :: Type a -> Some Type
baseType a@TCon{}    = Like a
baseType (TCon1 _ a) = baseType a
baseType (a :-> b)   = baseType a

shrinkClosed :: Signature sig => Type a -> Closed sig a -> [Closed sig a]
shrinkClosed a (Closed v) = do
  v <- shrinkTerm it a v
  Closed v : do v' <- shrinkTerm it a v; return (Closed v)

shrinkScheme :: Some Scheme -> [Some Scheme]
shrinkScheme (Like s) =
  case s of
    SCon c | isUnit c -> []
    SVarX     -> [Like unit]
    SVarY     -> [Like unit, Like SVarX]
    SCon c    -> [Like unit, Like SVarX, Like SVarY]
    SCon1 c a -> Like a : [ Like (SCon1 c a) | Like a <- shrinkScheme (Like a) ]
    a :=> b   -> [Like a, Like b] ++
                 [ Like (a :=> b) | Like a <- shrinkScheme (Like a) ] ++
                 [ Like (a :=> b) | Like b <- shrinkScheme (Like b) ]
  where
    unit = it :: Scheme (SCON "Unit")
    isUnit c = not $ null $ c =?= (it :: Name "Unit")

shrinkType :: Type a -> [Some Type]
shrinkType a = shrinkSomeType (Like a)

shrinkSomeType :: Some Type -> [Some Type]
shrinkSomeType (Like s) =
  case s of
    TCon c    -> [Like unit | not $ isUnit c]
    TCon1 c a -> Like a : [ Like (TCon1 c a) | Like a <- shrinkType a ]
    a :-> b   -> [Like a, Like b] ++
                 [ Like (a :-> b) | Like a <- shrinkType a ] ++
                 [ Like (a :-> b) | Like b <- shrinkType b ]
  where
    unit = it :: Type (TCON "Unit")
    isUnit :: Name c -> Bool
    isUnit c = not $ null $ c =?= (it :: Name "Unit")

-- Some utility functions --

spineLen :: SpineS g a -> Int
spineLen HeadS{}  = 0
spineLen (ArgS s _) = 1 + spineLen s

wkSpine :: Spine g a -> Spine (g ::> b) a
wkSpine (Head x a b)  = Head (Vs x) a b
wkSpine (Arg s a) = wkSpine s `Arg` a

wkSpineS :: SpineS g a -> SpineS (g ::> b) a
wkSpineS (HeadS x i j)  = HeadS (Vs x) i j
wkSpineS (ArgS s a) = wkSpineS s `ArgS` a

instSpine :: Cxt Scheme g -> SpineS g s -> Type x -> Type y -> Spine g (Inst s x y)
instSpine cxt (HeadS v x y) i j = lemInst (cxt ! v) x y i j ~$ Head v (inst x i j) (inst y i j)
instSpine cxt (ArgS sp s) x y = Arg (instSpine cxt sp x y) (inst s x y)

instSpineS :: Cxt Scheme g -> SpineS g s -> Scheme x -> Scheme y -> SpineS g (InstS s x y)
instSpineS cxt (HeadS v x y) i j = lemInstS (cxt ! v) x y i j ~$ HeadS v (instS x i j) (instS y i j)
instSpineS cxt (ArgS sp a) i j = ArgS (instSpineS cxt sp i j) (instS a i j)

typeOf :: Cxt Scheme g -> Term g a -> Type a
typeOf cxt v = case v of
  Var x a b -> inst (cxt ! x) a b
  f :@ v    -> case typeOf cxt f of _ :-> b -> b
  Lam a v   -> a :-> typeOf (cxt :> mono a) v


-- (===) :: (Show a, Eq a) => a -> a -> Property
-- x === y = whenFail (putStrLn $ show x ++ " /= " ++ show y) (x == y)

cheat = setStdGen (mkStdGen 27)

instance Arbitrary (Some Scheme) where
  arbitrary = sized $ genScheme (tyConGen (Proxy :: Proxy (MonadSig (State Int))))
  shrink = shrinkScheme

instance Arbitrary (Some Type) where
  arbitrary = sized $ genType (tyConGen (Proxy :: Proxy (MonadSig (State Int))))
  shrink = shrinkSomeType

prop_unify (Like s) (Like x) (Like y) =
  1 == length (unify s (mono $ inst s x y))

-- Pretty printing --------------------------------------------------------

class Pretty a where
  pretty :: a -> Doc
  prettyPrec :: Int -> a -> Doc

  pretty = prettyPrec 0
  prettyPrec _ = pretty

class Pretty1 (f :: k -> *) where
  pretty1     :: f a -> Doc
  prettyPrec1 :: Int -> f a -> Doc

  pretty1 = prettyPrec1 0
  prettyPrec1 _ = pretty1

instance Pretty1 f => Pretty (Some f) where
  prettyPrec p (Like x) = prettyPrec1 p x

instance Pretty1 f => Show (Some f) where showsPrec p = shows . prettyPrec p

instance Show (Closed sig a) where showsPrec p = shows . prettyPrec p

instance Pretty (Closed sig a) where
  prettyPrec p (Closed v) = prettyTerm ns p v
    where
      ns = mapCxt (K . cName) (sigEnv (Proxy :: Proxy sig))

instance Pretty (Glued sig a) where prettyPrec p (v :~ _) = prettyPrec p v
instance Show (Glued sig a)   where showsPrec p = shows . prettyPrec p
instance Pretty1 (Glued sig)  where prettyPrec1 = prettyPrec

instance Pretty1 ty => Pretty (Cxt ty g) where
  prettyPrec p Nil        = text "ε"
  prettyPrec p (Nil :> a) = prettyPrec1 p a
  prettyPrec p (cxt :> a) = parensIf (p > 5) $
    sep [ prettyPrec 5 cxt <> text ","
        , prettyPrec1 6 a ]

instance Pretty1 ty => Show (Cxt ty g) where
  showsPrec p = shows . prettyPrec p

deriving instance Show (Var g a)

parensIf :: Bool -> Doc -> Doc
parensIf True  = parens
parensIf False = id

prettyOp :: (Pretty a, Pretty b) => Bool -> Int -> Int -> String -> a -> b -> Doc
prettyOp p l r op x y = parensIf p $
  sep [ prettyPrec l x <+> pop, nest 2 $ prettyPrec r y ]
  where
    pop | op == ""  = P.empty
        | otherwise = text op

infixLeft, infixRight :: (Pretty a, Pretty b) => Int -> Int -> String -> a -> b -> Doc
infixLeft  p r = prettyOp (p > r) r (r + 1)
infixRight p r = prettyOp (p > r) (r + 1) r

instance Pretty (Var g a) where
  pretty x = text (show x)

instance Pretty1 (Var g) where pretty1 = pretty

instance Pretty (Int -> Doc) where
  prettyPrec p f = f p

data K a b = K { unK :: a }

instance Pretty1 Type   where prettyPrec1 = prettyPrec
instance Pretty1 Scheme where prettyPrec1 = prettyPrec

instance Pretty (Type a) where
  prettyPrec p a = case a of
    TCon c    -> pretty c
    TCon1 c a | "[]" == show c -> brackets (pretty a)
    TCon1 c a -> parensIf (p > 9) $ sep [ pretty c, nest 2 $ prettyPrec 10 a ]
    a :-> b   -> parensIf (p > 4) $ sep [ prettyPrec 5 a <+> text "->", prettyPrec 4 b ]

instance Pretty (Name s) where
  pretty (Name x) = text $ symbolVal x
instance Show (Name s) where show = show . pretty

instance Show (Scheme s) where showsPrec p = shows . prettyPrec p

instance Pretty (Scheme s) where
  prettyPrec p s =
    parensIf (p > 0 && not (null forallStr)) $
    sep [ text forallStr
        , nest 2 $ prettyPrec p' (inst s (it :: Type (TCON "x")) (it :: Type (TCON "y"))) ]
    where
      useX = null $ occursCheck SVarX s
      useY = null $ occursCheck SVarY s
      p' | null forallStr = p
         | otherwise      = 0
      forallStr = case (useX, useY) of
        (False, False) -> ""
        (True,  False) -> "forall x."
        (False, True)  -> "forall y."
        (True,  True)  -> "forall x y."

prettyTerm :: Cxt (K String) g -> Int -> Term g a -> Doc
prettyTerm env = pr names env
  where
    names = [ s ++ [c] | s <- "" : names, c <- ['a'..'z'] ]

    v2i :: forall g a. Var g a -> Int
    v2i Vz     = 0
    v2i (Vs x) = 1 + v2i x

    varName :: forall g a. Var g a -> Cxt (K String) g -> String
    varName x env = unK $ env ! x

    pr :: forall g a. [String] -> Cxt (K String) g -> Int -> Term g a -> Doc
    pr ns env p a =
      case a of
        Var x _ _ :@ u :@ v | isInfix s ->
            (either (infixLeft p) (infixRight p) fx) s (flip (pr ns env) u) (flip (pr ns env) v)
          where s = varName x env
                fx = fixity s
        Var x _ _ :@ u | isInfix s ->
          parens $ pr ns env fx u <+> text s
          where s  = varName x env
                fx = either id (+ 1) $ fixity s
        Var x _ _ | isInfix s -> parens (text s)
          where s = varName x env
        Var x a b -> text (varName x env)
                     -- <+> sep [ text "@" <+> prettyPrec 9 a
                     --         , text "@" <+> prettyPrec 9 b ]
        f :@ x    -> infixLeft p 10 "" (flip (pr ns env) f) (flip (pr ns env) x)
        Lam _ _   ->
          case prettyLam ns env a of
            (xs, d) -> parensIf (p > 0) $
              sep [ text "\\" <> fsep xs <+> text "->"
                  , nest 2 d ]
          where
            annotate :: forall a. String -> Type a -> Doc
            annotate x _ = text x
            -- annotate x t = parens $ text x <+> text "::" <+> pretty t

            prettyLam :: forall g a. [String] -> Cxt (K String) g -> Term g a -> ([Doc], Doc)
            prettyLam (x:ns) env (Lam a e) =
              (annotate x a :) *** id $ prettyLam ns (env :> K x) e
            prettyLam ns env e = ([], pr ns env 0 e)

isInfix (c:_) = not $ isAlphaNum c || elem c "[("

fixity ":"   = Right 5
fixity ">>=" = Left 1
fixity "+"   = Left 6
fixity _     = Left 9

class    IsCxt cxt where cxtShape :: Cxt (K ()) cxt
instance IsCxt NIL where cxtShape = Nil
instance IsCxt g => IsCxt (g ::> a) where cxtShape = cxtShape :> K ()

defaultNames :: IsCxt g => Cxt (K String) g
defaultNames = names cxtShape 0
  where
    names :: Cxt a g -> Int -> Cxt (K String) g
    names Nil _ = Nil
    names (g :> _) i = names g (i + 1) :> K ("X" ++ show i)

instance IsCxt g => Pretty1 (Term g) where
  prettyPrec1 = prettyPrec

instance IsCxt g => Pretty (Term g a) where
  prettyPrec = prettyTerm defaultNames

instance Show (Type a) where showsPrec p = shows . prettyPrec p
instance IsCxt g => Show (Term g a)  where showsPrec p = shows . prettyPrec p

instance Pretty (a :=: b) where
  pretty = text . show

deriving instance Show (a :=: b)

deriving instance Show (Spine g a)

deriving instance Show (Unified a b)