content: {-----------------------------------------------------------------------------
  An Effects-and-Handlers Implementation of Hindley-Milner Typechecking
  
  Points of interest include
  
  1. I use a free monad indexed over "Time" to represent computations which
     take time, during which progress can occur. Progress is packaged as the
     morphisms of the Time category, and is propagated rather quietly.
  2. Instead of representing the typing context explicitly as a data structure,
     all contextualisation is done by effect handlers, responding to the effects
     enabled by the monad.
  
  -----------------------------------------------------------------------------}
  
  
  {-# LANGUAGE DataKinds, GADTs, KindSignatures, RankNTypes, StandaloneDeriving,
      QuantifiedConstraints, LambdaCase, ScopedTypeVariables, TypeOperators,
      TypeFamilies, UndecidableInstances, ConstraintKinds, TupleSections,
      IncoherentInstances, OverlappingInstances, PatternSynonyms,
      ViewPatterns, TypeOperators, LiberalTypeSynonyms #-}
  
  module EffW where
  
  import Data.Kind      -- it's going to be this sort of adventure
  import Unsafe.Coerce  -- and this sort
  
  ------------------------------------------------------------------------------
  -- Indexing
  ------------------------------------------------------------------------------
  
  -- index-respecting functions
  
  type (-:>) s t i = s i -> t i
  type All f = forall i. f i
  infixr 3 -:>
  
  
  data (:*) (s :: a -> Type)(t :: a -> Type)(i :: a) where
    (:*) :: s i -> t i -> (s :* t) i
  infixr 4 :*
  
  newtype K (t :: Type)(i :: a) = K {unK :: t} deriving Eq
  
  kmap :: (s -> t) -> All (K s -:> K t)
  kmap f (K s) = K (f s)
  
  data Some (f :: a -> Type) :: Type where
    Some :: f x -> Some f
  deriving instance (forall x. Show (f x)) => Show (Some f)
  
  
  ------------------------------------------------------------------------------
  -- Time
  ------------------------------------------------------------------------------
  
  data Time
  
  -- Step sequences through time are snoc-lists which join up like dominoes
  data (&>) (s :: Time)(t :: Time) :: Type where
    Now  :: t &> t
    (:<) :: r &> s -> Step s t -> r &> t
  
  class Timed (v :: Time -> Type) where
    -- value-level &> gives the action of type-level &> ...
    (&>) :: v s -> s &> t -> v t
    -- ...by iterating the 
    step :: v s -> Step s t -> v t
    v &> Now = v
    v &> (sz :< s) = step (v &> sz) s
  
  instance Timed ((&>) r) where
    step = (:<)
  
  now :: v s -> s &> s
  now _ = Now
  
  instance Timed (K a) where
    K a &> _ = K a
    step (K a) _ = K a
  
  instance (Timed s, Timed t) => Timed (s :* t) where
    step (s :* t) u = step s u :* step t u
  
  
  class LE (s :: Time)(t :: Time) where
    lesson :: s &> t
  
  class SegTo (t :: Time) where
    type SegFrom t :: Time
    seg :: SegFrom t &> t
  
  instance LE s s where
    lesson = Now
  
  instance (LE s (SegFrom t), SegTo t) => LE s t where
    lesson = lesson &> seg
  
  type Kripke v t = forall u. LE t u => v u
  
  kripke :: Timed v => All (v -:> Kripke v)
  kripke v = v &> lesson
  
  kripkefy :: Timed v => All (v -:> (Kripke v -:> f) -:> f)
  kripkefy v k = k (kripke v)
  
  
  class MoTime (m :: (Time -> Type) -> (Time -> Type)) where
    retNow :: Timed v => All (v -:> m v)
    (>>>=)  :: (Timed f, Timed g)
           => m f s
           -> (forall t. (SegTo t, SegFrom t ~ s) =>
               Kripke f t -> m g t)
           -> m g s
    (>>>)  :: (Timed f, Timed g)
           => m f s
           -> (forall t. (SegTo t, SegFrom t ~ s) => m g t)
           -> m g s
    mf >>> mg = mf >>>= \ _ -> mg
  
  split
    :: (Timed s, Timed t)
    => All ((s :* t) -:> (Kripke s -:> Kripke t -:> f) -:> f)
  split (s :* t) f = f (kripke s) (kripke t)
  
  
  ------------------------------------------------------------------------------
  -- de Bruijn Indices
  ------------------------------------------------------------------------------
  
  -- natural numbers, for how many things are in scope
  data N = Z | S N
  
  -- de Bruijn indices, for which things are in scope
  data Inx (n :: N) :: Type where
    Zi :: Inx (S n)
    Si :: Inx n -> Inx (S n)
  deriving instance Eq (Inx n)
  instance Show (Inx n) where
    show = show . go where
      go :: forall n. Inx n -> Int
      go Zi     = 0
      go (Si i) = 1  + go i
  
  -------
  -- if you know an Inx n, n isn't Z
  data Pos (n :: N) :: Type where
    Pos :: Pos (S n)
  
  inxPos :: Inx n -> Pos n
  inxPos Zi = Pos
  inxPos (Si _) = Pos
  -------
  
  -- thin j i is never j
  thin :: Inx (S n) -> Inx n -> Inx (S n)
  thin  Zi        i  = Si i
  thin (Si j)  Zi    = Zi
  thin (Si j) (Si i) = Si (thin j i)
  
  -- thick j j = Nothing; thick j (thin j i) = i
  thick :: Inx (S n) -> Inx (S n) -> Maybe (Inx n)
  thick  Zi     Zi    = Nothing
  thick  Zi    (Si j) = pure j
  thick (Si j)  Zi    = case inxPos j of Pos -> pure Zi
  thick (Si j) (Si i) = case inxPos j of Pos -> Si <$> thick i j
  
  
  ------------------------------------------------------------------------------
  -- What changes with time?
  ------------------------------------------------------------------------------
  
  data Step (s :: Time)(t :: Time) :: Type where
    (:/:) :: Ty Z -> ExVar -> Step t t
  
  type ExVar =
    ( String  -- for mnemonic purposes
    , Int     -- for freshness
    )
  
  data Ty (n :: N)
    = E ExVar    -- existential type variables
    | U (Inx n)
    | Ty n :-> Ty n
    deriving (Show, Eq)
  
  closed :: Ty Z -> Ty n
  closed = unsafeCoerce
  
  newtype Tyme (t :: Time) = Ty (Ty Z) deriving (Show, Eq)
  
  tiE :: All (K ExVar -:> Tyme)
  tiE (K e) = Ty (E e)
  
  data UnArrEh (i :: Time) where
    UnArrAye :: Kripke Tyme i -> Kripke Tyme i -> UnArrEh i
    UnArrNaw :: UnArrEh i
  
  unArrEh :: Tyme i -> UnArrEh i
  unArrEh (Ty (s :-> t)) = kripkefy (Ty s) $ \ s -> kripkefy (Ty t) $ \ t ->
    UnArrAye s t
  unArrEh _ = UnArrNaw
  
  pattern (:-&>) :: Kripke Tyme i -> Kripke Tyme i -> Tyme i
  pattern s :-&> t <- (unArrEh -> UnArrAye s t)
  
  (-&>) :: All (Tyme -:> Tyme -:> Tyme)
  Ty a -&> Ty b = Ty (a :-> b)
  
  instance Timed Tyme where
    step (Ty t) (r :/: x) = Ty (subst r x t)
  
  subst :: Ty Z -> ExVar -> Ty n -> Ty n
  subst r x = go where
    go (E y) | x == y = closed r
    go (s :-> t) = go s :-> go t
    go t = t
  
  data Sch (n :: N)
    = T (Ty n)
    | P (Sch (S n))
    deriving Show
  
  newtype Schime (t :: Time) = Sch (Sch Z) deriving Show
  
  monotype :: All (Tyme -:> Schime)
  monotype (Ty t) = Sch (T t)
  
  dep :: ExVar -> Ty n -> Bool
  dep x (E y) = x == y
  dep x (s :-> t) = dep x s || dep x t
  dep _ _ = False
  
  instance Timed Schime where
    step (Sch s) (r :/: x) = Sch (go s) where
      go :: All (Sch -:> Sch)
      go (T t) = T (subst r x t)
      go (P s) = P (go s)
      
  stan :: Sch (S Z) -> All (Tyme -:> Schime)
  stan s (Ty r) = Sch (sub Zi s) where
    sub :: Inx (S n) -> Sch (S n) -> Sch n
    sub j (T t) = T (go t) where
      go (s :-> t) = go s :-> go t
      go (U i) = case thick j i of
        Nothing -> closed r
        Just i  -> U i
      go (E e) = E e
    sub j (P p) = P (sub (Si j) p)
  
  gen :: All (K ExVar -:> Schime -:> Schime)
  gen (K e) (Sch s) = case go Zi s of (s', b) -> if b then Sch (P s') else Sch s
   where
  
    go :: Inx (S n) -> Sch n -> (Sch (S n), Bool)
    go j (T t) = case euTy e j t of (t', b) -> (T t', b)
    go j (P s) = case go (Si j) s of (s', b) -> (P s', b)
  
    euTy :: ExVar -> Inx (S n) -> Ty n -> (Ty (S n), Bool)
    euTy e j (E x) = if e == x then (U j, True) else (E x, False)
    euTy _ j (U i) = (U (thin j i), False)
    euTy e j (s :-> t) = case (euTy e j s, euTy e j t) of
      ((s, a), (t, b)) -> (s :-> t, a || b)
  
  
  ------------------------------------------------------------------------------
  -- Free Timed monads
  ------------------------------------------------------------------------------
  
  data TiMo
    (c :: (Time -> Type) -> Time -> Type)
    (v :: Time -> Type)
    (s :: Time)
    :: Type where
    RetNow :: v s -> TiMo c v s
    Call :: forall c r v s
          . c r s
         -> All ((&>) s -:> r -:> TiMo c v)
         -> TiMo c v s
  
  instance (forall r. Timed (c r), Timed v) => Timed (TiMo c v) where
    RetNow v &> u = RetNow (v &> u)
    Call c k &> u = Call (c &> u) $ \ w r -> k (u &> w) r
    step (RetNow v) u = RetNow (step v u)
    step (Call c k) u = Call (step c u) $ \ w r -> k ((Now :< u) &> w) r
  
  instance MoTime (TiMo c) where
    retNow = RetNow
    (>>>=) :: forall c f g s. (Timed f, Timed g)
           => TiMo c f s
           -> (forall t. (SegTo t, SegFrom t ~ s) =>
               Kripke f t -> TiMo c g t)
           -> TiMo c g s
    RetNow v >>>= k = leap (now v) (k (kripke v))
    Call c j >>>= k = Call c $ \ u r ->
      j u r >>>= jump u seg
      where
        jump :: forall t t'. s &> t -> t &> t'
             -> Kripke f t' -> TiMo c g t'
        jump u w f = leap (u &> w) (k f)
  
  op :: Timed r => c r s -> TiMo c r s
  op c = Call c $ \ _ r -> RetNow r
  
  leap :: forall s t x
        . s &> t
       -> ((SegTo t, SegFrom t ~ s) => x)
       -> x
  leap u k = case mkStep u of
    Dict -> k
  
  ------------------------------------------------------------------------------
  -- Seek not to know how the sausage is cooked
  ------------------------------------------------------------------------------
  
  data Dict (c :: Constraint) :: Type where
    Dict :: c => Dict c
  
  data FakeLe s t = FakeLe (s &> t)
  
  mkStep :: forall s t. s &> t -> Dict (SegTo t, (SegFrom t ~ s))
  mkStep u = case (foo, baz) of
      (Dict, Dict) -> Dict
    where
    foo :: Dict (SegTo t)
    foo = unsafeCoerce (FakeLe u)
    bar :: Dict (s ~ s)
    bar = Dict
    baz :: Dict (SegFrom t ~ s)
    baz = unsafeCoerce bar
  
  
  ------------------------------------------------------------------------------
  -- Effects
  ------------------------------------------------------------------------------
  
  data W (r :: Time -> Type)(i :: Time) :: Type where
    Next -- next fresh number, please (handled by nonce)
      :: W (K Int) i
    VSch -- look up the program variable (handled by decl)...
      :: ProgVar      -- ...with this name, and...
      -> W Schime i  -- tell me its type scheme
    Inst -- instantiate (handled by bloc)...
      :: Schime i  -- ...this type scheme...
      -> W Tyme i  -- ...to this type (by guessing the type parameters)
    Make -- make a definition (handled by guessing)...
      :: [ExVar]  -- ...with these dependencies to be extruded...
      -> Tyme i   -- ...of this type...
      -> ExVar    -- ...for this variable
      -> W (K ()) i
    Barf :: W f i
  
  instance Timed (W r) where
    step Next u = Next
    step (VSch x) u = VSch x
    step (Inst s) u = Inst (step s u)
    step (Make ds t x) u = Make ds (step t u) x
    step Barf u = Barf
  
  
  ------------------------------------------------------------------------------
  -- Handlers
  ------------------------------------------------------------------------------
  
  -- handle Next...
  nonce :: Timed v => Int -> All (TiMo W v -:> TiMo W v)
  -- ...by giving the next number and rehandling the continuation with increment
  nonce n (Call Next k) = nonce (n + 1) $ k Now (K n)
  -- forward everything else
  nonce n (RetNow v) = RetNow v
  nonce n (Call c k) = Call c $ \ u r -> nonce n $ k u r
  
  -- pick a fresh existential variable using Next
  fresh :: String -> TiMo W (K ExVar) i
  fresh x =
    op Next >>>= \ i ->
    retNow (kmap (x,) i)
  
  type ProgVar = String
  -- handle VSch...
  decl :: Timed v => ProgVar -> All (Schime -:> TiMo W v -:> TiMo W v)
  -- ...by giving the scheme if we're looking up this decl
  decl x s (Call (VSch y) k) | x == y = decl x s $ k Now s
  -- forward everything else, but...
  decl x s (RetNow v) = RetNow v
  -- ...be sure to update the scheme in the light of progress
  decl x s (Call c k) = Call c $ \ u r -> decl x (s &> u) $ k u r
    
  
  -- handle Inst requests
  --   (this is fatsemi in Gundry-McBride-McKinna
  --   , doorstop in Krishnaswami-Dunfield)
  bloc :: All (TiMo W Tyme -:> TiMo W Schime)
  -- if we're instantiating a monotype, we're done
  bloc (Call (Inst (Sch (T t))) k) = bloc $ k Now (Ty t)
  -- if we're instantiating a polytime, we're inventing a fresh existential var
  -- and guessing it
  bloc (Call (Inst (Sch (P s))) k) =
    fresh "x" >>>= \ e ->
    guessing e $ bloc $
      op (Inst (stan s (tiE e))) >>>= \ t ->
      k lesson t
  -- retNow wraps the type
  bloc (RetNow (Ty t)) = RetNow (Sch (T t))
  -- otherwise forward
  bloc (Call c k) = Call c $ \ u r -> bloc $ k u r
  
  -- handle Make, but also do generalisation (note we're computing type schemes)
  guessing :: All (K ExVar -:> TiMo W Schime -:> TiMo W Schime)
  -- when Make shows up, we have four possibilities
  guessing (K e) (Call c@(Make ds (Ty t) x) k) = case (e == x, dep e t) of
    -- (is it me?, do I occur in the definiens)
    (True, True)  -- it's me and the occur check failed; oh noes!
      -> op Barf
    (True, False)  -- it's me, so extrude my dependencies and substitute me!
      -> foldr (guessing . K) (k (Now :< (t :/: x)) (K ())) ds
    (False, True)  -- it's not me, but I occur, so extrude me!
      -> Call (Make (e : ds) (Ty t) x) k
    (False, False)  -- it's nothing to do with me, so leave alone!
      -> Call c $ \ u r -> guessing (K e) $ k u r
  -- nobody made me; I could be anything; pawn becomes queen!
  guessing e (RetNow s) = RetNow (gen e s)
  -- forward the rest (the update is a no-op)
  guessing (K e) (Call c k) = Call c $ \ u r -> guessing (K e) $ k u r
  
  
  ------------------------------------------------------------------------------
  -- Run
  ------------------------------------------------------------------------------
  
  run :: TiMo W f i -> Maybe (Some f)
  run (RetNow r) = Just (Some r)
  run _ = Nothing
  
  
  ------------------------------------------------------------------------------
  -- Core ML
  ------------------------------------------------------------------------------
  
  data Tm
    = V ProgVar              -- program variables
    | ProgVar :=> Tm         -- lambda
    | Tm :$ Tm               -- application
    | (ProgVar, Tm) :/ Tm    -- let
    deriving Show
  
  infixr 3 :/
  infixr 4 :=>
  infixl 5 :$
  
  -- ensure that a type is a function type, giving back source and target
  funTy :: All (Tyme -:> TiMo W (Tyme :* Tyme))
  -- if it's already a function type, crack on!
  funTy (Ty (s :-> t)) = retNow (Ty s :* Ty t)
  -- otherwise, invent a function type and constrain
  funTy u = kripkefy u $ \ u ->
    op (Inst (Sch (P (P (T (U (Si Zi) :-> U Zi)))))) >>>= \ f ->
    unify u f >>>= \ _ ->
    funTy f
  
  -- guess a type
  guess :: All (TiMo W Tyme)
  guess = op (Inst (Sch (P (T (U Zi)))))
  
  -- make types the same!
  unify :: All (Tyme -:> Tyme -:> TiMo W (K ()))
  -- if they're already the same, we're done
  -- (note that catches trivial occur check failures, leaving only bad ones)
  unify a b | a == b = retNow (K ())
  -- rigid-rigid decomposition
  unify (s0 :-&> t0) (s1 :-&> t1) = unify s0 s1 >>> unify t0 t1
  -- flex problem? demand a definition!
  unify (Ty (E e)) t = op (Make [] t e)
  unify t (Ty (E e)) = op (Make [] t e)
  -- anything else is hopeless
  unify _ _ = op Barf
  
  infer :: Tm -> All (TiMo W Tyme)
  infer (V x) =
    op (VSch x) >>>= \ s ->   -- look up the declaration
    op (Inst s)              -- instantiate it
  infer (x :=> b) =
    guess >>>= \ s ->                           -- guess the domain
    decl x (monotype s) (infer b) >>>= \ t ->   -- declare x monomorphically
    retNow (s -&> t)                            -- assemble the function type
  infer (f :$ a) =
    infer f >>>= \ ft ->    -- infer the function's type
    infer a >>>= \ at ->    -- infer the argument's type
    funTy ft >>>= \ st ->   -- see the function's type as a function type
      split st $ \ s t ->   -- split into domain and codomain
        unify at s >>>      -- unify argument's type with domain
        retNow t            -- give back the codomain
  infer ((x, d) :/ b) =
    bloc (infer d) >>>= \ s ->  -- get a type scheme for the definiens
    decl x s (infer b)          -- declare the definiendum and infer the body
  
  
  ------------------------------------------------------------------------------
  -- Entry point
  ------------------------------------------------------------------------------
  
  whatIs :: Tm -> Maybe (Sch Z)
  whatIs t = case run . nonce 0 . bloc $ infer t of
    Just (Some (Sch s)) -> Just s
    _ -> Nothing
  
  skk :: Maybe (Sch Z)
  skk = whatIs $
    ("s", "f" :=> "a" :=> "g" :=> V "f" :$ V "g" :$ (V "a" :$ V "g")) :/
    ("k", "x" :=> "g" :=> V "x") :/
    V "s" :$ V "k" :$ V "k"
  