content: -- \| The untyped lambda calculus with pattern matching
  --
  
  -- |
  -- Module      : Pat
  -- Description : Untyped lambda calculus, with pattern matching
  -- Stability   : experimental
  --
  -- An implementation of the untyped lambda calculus with pattern matching.
  --
  -- This example extends the lambda calculus with constants (like 'nil and 'cons)
  -- and arbitrary pattern matching. Case expressions include a list of branches,
  -- where each branch is a pattern and a right-hand side. The pattern can bind
  -- multiple variables and the index ensures that the rhs matches the number of
  -- variables bound in the pattern.
  -- 
  -- This example also includes pairs and "irrefutable patterns" i.e. let binding
  -- that can deeply deconstruct cons pairs (only).
  module Pat where
  
  import Rebound
  
  import Rebound.Bind.PatN as PatN
  import qualified Rebound.Bind.Pat as Pat
  import Rebound.Bind.Scoped qualified as Scoped
  import Data.Maybe qualified as Maybe
  import Data.Type.Equality
  import Data.Fin ( Fin, f0, f1 )
  import Data.Fin qualified as Fin
  import Data.Vec qualified as Vec
  
  ----------------------------------------------
  
  -- * Syntax
  
  ----------------------------------------------
  
  -- The untyped lambda calculus extended with
  -- symbols ("con"stants) and pattern matching
  -- expression (case)
  -- A constant applied to any number of arguments
  -- is a value
  data Exp (n :: Nat) where
    Var :: Fin n -> Exp n
    Lam :: Bind1 Exp Exp n -> Exp n
    App :: Exp n -> Exp n -> Exp n
    LetPair  :: Exp n -> Branch PairPat n -> Exp n
    -- ^ deep pattern matching against pairs (only)
    Con :: String -> Exp n
    -- ^ constant (or symbol) like 'cons or 'nil'
    Case :: Exp n -> [Branch Pat n] -> Exp n
    -- ^ deep pattern matching against any pattern
     
  
  -- Each branch in a case expression is a pattern binding,
  -- i.e. a data structure that binds m variables in some
  -- expression body with scope n
  -- Here, the variable m does not appear
  -- in the result type `Branch pat n`, so is an existential.
  data Branch pat (n :: Nat) where
    Branch :: SNatI m => Pat.Bind Exp Exp (pat m) n -> Branch pat n
  
  -- Patterns for case expressions.
  -- The index `m` in the pattern is the number of occurrences of
  -- PVar, i.e. the number variables bound by the pattern.
  -- These variables are ordered left to right.
  -- For example (PCon "cons" `PApp` PVar `PApp` PVar) is the
  -- representation of the pattern "cons x y", which binds two
  -- variables.
  -- To prevent patterns of the form "x y z", this type is split
  -- into top level patterns (Pat) and applications of constants (ConApp)
  data Pat (m :: Nat) where
    PVar :: Pat N1 -- binds exactly one variable
    PHead :: ConApp m -> Pat m
  
  data ConApp (m :: Nat) where
    PCon :: String -> ConApp N0 -- binds zero variables
    PApp :: ConApp m1 -> Pat m2 -> ConApp (m2 + m1)
  
  -- Patterns for pairs only, a special case of the above
  data PairPat (m :: Nat) where
    PPVar :: PairPat N1
    PPair :: PairPat m1 -> PairPat m2 -> PairPat (m2 + m1)
  
  ----------------------------------------------
  
  -- * Sized instance
  
  ----------------------------------------------
  
  -- Any type that is used as a pattern must be an
  -- instance of the `Sized` type class, so that the library
  -- can determine the number of binding variables both
  -- statically and dynamically.
  
  -- The `Pat` type tells us how many variables are bound
  -- the pattern with the index `n`. We can also recover
  -- that number from the pattern itself by counting the number
  -- of occurrences of `PVar`.
  
  instance Sized (Pat m) where
    type Size (Pat m) = m
  
    size :: Pat m -> SNat (Size (Pat m))
    size PVar = s1
    size (PHead p) = size p
  
  
  instance Sized (ConApp m) where
    type Size (ConApp m) = m
  
    size :: ConApp m -> SNat (Size (ConApp m))
    size (PApp p1 p2) = sPlus (size p2) (size p1)
    size (PCon s) = s0
  
  
  instance Sized (PairPat m) where
    type Size (PairPat m) = m
  
    size :: PairPat m -> SNat (Size (PairPat m))
    size PPVar = s1
    size (PPair p1 p2) = sPlus (size p2) (size p1)
  
  
  ----------------------------------------------
  
  -- * Substitution
  
  ----------------------------------------------
  
  instance SubstVar Exp where
    var :: Fin n -> Exp n
    var = Var
  
  instance Shiftable Exp where
    shift = shiftFromApplyE @Exp
  
  instance Subst Exp Exp where
    applyE :: Env Exp n m -> Exp n -> Exp m
    applyE r (Var x) = applyEnv r x
    applyE r (Lam b) = Lam (applyE r b)
    applyE r (App e1 e2) = App (applyE r e1) (applyE r e2)
    applyE r (Con s) = Con s
    applyE r (Case e brs) = Case (applyE r e) (map (applyE r) brs)
    applyE r (LetPair e1 b) = LetPair (applyE r e1) (applyE r b)
  
  
  instance Shiftable (Branch pat) where
    shift = shiftFromApplyE @Exp
  
  instance Subst Exp (Branch pat) where
    applyE :: Env Exp n m -> Branch pat n -> Branch pat m
    applyE r (Branch bnd) = Branch (applyE r bnd)
  
  
  ----------------------------------------------
  -- Example terms
  ----------------------------------------------
  
  -- The identity function "λ x. x". With de Bruijn indices
  -- we write it as "λ. 0"
  t0 :: Exp Z
  t0 = Lam (bind1 (Var f0))
  
  -- A larger term "λ x. λy. x (λ z. z z)"
  -- λ. λ. 1 (λ. 0 0)
  t1 :: Exp Z
  t1 =
    Lam
      ( bind1
          ( Lam
              ( bind1
                  ( Var f1
                      `App` (Lam (bind1 (Var f0)) `App` Var f0)
                  )
              )
          )
      )
  
  -- "head function"
  -- \x -> case x of [nil -> x ; cons y z -> y]
  t2 :: Exp Z
  t2 =
    Lam
      ( bind1
          ( Case
              (Var f0)
              [ Branch
                  ( Pat.bind @(Pat N0)
                      (PHead (PCon "Nil"))
                      (Var f0)
                  ),
                Branch
                  ( Pat.bind @(Pat N2)
                      (PHead (PCon "Cons" `PApp` PVar `PApp` PVar))
                      (Var f0)
                  )
              ]
          )
      )
  
  -- a "list"  ['a','b']
  t3 :: Exp Z
  t3 = Con "cons" `App` Con "a" `App` (Con "cons" `App` Con "b" `App` Con "nil")
  
  --------------------------------------------------------------
  
  -- * Show implementation
  
  --------------------------------------------------------------
  
  -- >>> t0
  -- λ. 0
  
  -- >>> t1
  -- λ. λ. 1 (λ. 0 0)
  
  -- >>> t2
  -- λ. case 0 of [Nil => 0,(Cons V) V => 0]
  
  -- >>> t3
  -- (cons a) ((cons b) nil)
  
  
  instance Show (Exp n) where
    showsPrec :: Int -> Exp n -> String -> String
    showsPrec _ (Var x) = shows x
    showsPrec d (App e1 e2) =
      showParen (d > 0) $
        showsPrec 10 e1
          . showString " "
          . showsPrec 11 e2
    showsPrec d (Lam b) =
      showParen (d > 10) $
        showString "λ. "
          . shows (getBody1 b)
    showsPrec d (Con s) = showString s
    showsPrec d (Case e brs) =
      showParen (d > 10) $
        showString "case "
          . shows e
          . showString " of "
          . shows brs
    showsPrec d (LetPair e (Branch b)) = 
      showString "let "
      . shows (Pat.getPat b)
      . showString " = "
      . shows e
      . showString " in "
      . showsPrec d (Pat.getBody b)
  
  instance Show (PairPat m) where
    showsPrec :: Int -> PairPat m -> String -> String
    showsPrec d PPVar = showString "V"
    showsPrec d (PPair p1 p2) =
      showParen True $
        shows p1
          . showString ","
          . shows p2
  
  instance Show (Pat m) where
    showsPrec :: Int -> Pat m -> String -> String
    showsPrec d PVar = showString "V"
    showsPrec d (PHead p) = showsPrec d p
  
  instance Show (ConApp m) where
    showsPrec d (PApp p1 p2) =
      showParen (d > 0) $
        showsPrec 10 p1
          . showString " "
          . showsPrec 11 p2
    showsPrec d (PCon s) = showString s
  
  
  -- In a `PatBind` term, we can access the pattern with `getPat`
  -- and the RHS with `getBody`
  instance Show (Branch Pat n) where
    showsPrec :: Int -> Branch Pat n -> String -> String
    showsPrec d (Branch bnd) =
      shows (Pat.getPat bnd)
        . showString " => "
        . showsPrec d (Pat.getBody bnd)
  
  --------------------------------------------------------------
  
  -- * Eq implementation
  
  --------------------------------------------------------------
  
  -- We would like to derive equality for patterns, i.e. 
  -- 
  --     deriving instance (Eq (Pat m))
  -- 
  -- but because of the application case, this process fails.
  -- We don't know that each subpattern binds the same
  -- number of variables!
  
  
  -- Therefore to compare Pats for equality, we generalize the
  -- `testEquality` function from Data.Type.Equality. (This
  -- class is often used for comparisons between indexed types.
  -- but only works if the index is the last type parameter.
  -- In our case, we need to produce an equality for the
  -- first type parameter.)
  -- This function can be applied, even if the number of
  -- pattern-bound variables are not known to be equal.
  -- (cf. m1 and m2 below). If the patterns are indeed equal,
  -- then `patEq` *also* returns a proof that the indices
  -- are equal. (The type `a :~: b` is a GADT with a single
  -- constructor `Refl` that can only be used when a and be are
  -- equal. Pattern matching on this GADT brings an equality
  -- between a and b into the context of the term.)
  
  instance PatEq (Pat m1) (Pat m2) where
    patEq PVar PVar = Just Refl
    patEq (PHead p1) (PHead p2) = do
      Refl <- patEq p1 p2
      return Refl
    patEq _ _ = Nothing
  
  instance PatEq (ConApp m1) (ConApp m2) where
    patEq (PApp p1 p2) (PApp p1' p2') = do
      Refl <- patEq p1 p1'
      Refl <- patEq p2 p2'
      return Refl
    patEq (PCon s1) (PCon s2) | s1 == s2 = Just Refl
    patEq _ _ = Nothing
  
  instance PatEq (PairPat m1) (PairPat m2) where
    patEq (PPair p1 p2) (PPair p1' p2') = do
      Refl <- patEq p1 p1'
      Refl <- patEq p2 p2'
      return Refl
    patEq PPVar PPVar = Just Refl
    patEq _ _ = Nothing
  
  
  -- the generalized equality can be used for the usual equality
  instance Eq (Pat m) where
    p1 == p2 = Maybe.isJust (patEq p1 p2)
  
  instance Eq (PairPat m) where
    p1 == p2 = Maybe.isJust (patEq p1 p2)
  
  instance SizeIndex PairPat p
  instance SizeIndex Pat p
  
  
  -- Because the Branch type is parameterized by a pattern type, `pat` of kind 
  -- `Nat -> Type` we need to make some assumptions about that type to construct
  -- the `Eq` instance. (1) we need to be able to test patterns for equality
  -- no matter what their size is. (2) we need to know that the index *is* the 
  -- size of the pattern, i.e. Size (pat m) ~ m. The `SizeIndex` class captures
  -- this relationship in a way that can be quantifed over all m.
  -- If we did not parameterize the `Branch` type by the pattern type, we would not
  -- need this complexity.
  
  instance (forall m. Eq (pat m),                    -- 1
            forall m. SizeIndex pat m)               -- 2
            => Eq (Branch pat n) where
    (==) :: Branch pat n -> Branch pat n -> Bool
    (Branch (p1 :: Pat.Bind Exp Exp (pat m1) n))
      == (Branch (p2 :: Pat.Bind Exp Exp (pat m2) n)) =
        case testEquality
          (size (Pat.getPat p1) :: SNat m1)
          (size (Pat.getPat p2) :: SNat m2) of
          Just Refl -> p1 == p2
          Nothing -> False
  
  
  -- With the instance above the derivable equality instance
  -- is alpha-equivalence
  deriving instance (Eq (Exp n))
  
  
  --------------------------------------------------------
  -- Pattern matching code
  --------------------------------------------------------
  
  p1 :: Pat N2
  p1 = PHead $ PApp (PApp (PCon "C") PVar) PVar
  
  p2 :: Pat N2
  p2 = PHead $ PApp (PApp (PCon "D") PVar) PVar
  
  e1 :: Exp N0
  e1 = App (App (Con "C") (Con "A")) (Con "B")
  
  e2 :: Exp N0
  e2 = App (App (Con "D") (Con "A")) (Con "C")
  
  -- >>> patternMatch p1 e1
  -- Just [(0,B),(1,A)]
  
  -- >>> patternMatch p2 e1
  -- Nothing
  
  -- >>> patternMatch p1 e2
  -- Nothing
  
  -- >>> patternMatch p2 e2
  -- Just [(0,C),(1,A)]
  
  
  -- | Compare a "pair" pattern with a pair pattern, potentially
  -- producing a substitution for all of the variables bound in the pattern.
  ppatternMatch :: PairPat p -> Exp m -> Maybe (Env Exp p m)
  ppatternMatch PPVar e = Just $ oneE e
  ppatternMatch (PPair p1 p2) (App (App (Con "cons") e1) e2) = do
    env1 <- ppatternMatch p1 e1
    env2 <- ppatternMatch p2 e2
    withSNat (size p2) $ return (env2 .++ env1)
  ppatternMatch _ _ = Nothing
  
  -- | Compare a pattern with an expression, potentially
  -- producing a substitution for all of the variables bound in the pattern.
  patternMatch :: Pat p -> Exp m -> Maybe (Env Exp p m)
  patternMatch PVar e = Just $ oneE e
  patternMatch (PHead p) e = patternMatchApp p e
  
  patternMatchApp :: ConApp p -> Exp m -> Maybe (Env Exp p m)
  patternMatchApp (PApp p1 p2) (App e1 e2) = do
    env1 <- patternMatchApp p1 e1
    env2 <- patternMatch p2 e2
    withSNat (size p2) $ return (env2 .++ env1)
  patternMatchApp (PCon s1) (Con s2) =
    if s1 == s2 then Just zeroE else Nothing
  patternMatchApp _ _ = Nothing
  
  -- Compare the scrutinee against multiple patterns and return 
  -- the matching branch
  findBranch :: Exp n -> [Branch Pat n] -> Maybe (Exp n)
  findBranch e [] = Nothing
  findBranch e (Branch bind : brs) =
    case patternMatch (Pat.getPat bind) e of
      Just r -> Just $ Pat.instantiate bind r
      Nothing -> findBranch e brs
  
  
  
  
  --------------------------------------------------------
  -- Eval and step
  --------------------------------------------------------
  
  {- We can write the usual operations for evaluating
     lambda terms to values -}
  -- big-step evaluation
  -- >>> eval t1
  -- λ. λ. 1 (λ. 0 0)
  -- >>> eval (t1 `App` t0)
  -- λ. λ. 0 (λ. 0 0)
  t4 = t2 `App` t3
  
  -- >>> t4
  -- λ. case 0 of [Nil => 0,(Cons V) V => 0] ((cons a) ((cons b) nil))
  -- >>> eval t4
  -- case (cons a) ((cons b) nil) of [Nil => (cons a) ((cons b) nil),(Cons V) V => 0]
  eval :: Exp n -> Exp n
  eval (Var x) = Var x
  eval (Lam b) = Lam b
  eval (App e1 e2) =
    let v = eval e2
     in case eval e1 of
          Lam b -> eval (instantiate1 b v)
          t -> App t v -- if cannot reduce, return neutral term
  eval (Con s) = Con s
  eval (Case e brs) =
    let v = eval e
     in case findBranch v brs of
          Just br -> eval br
          Nothing -> Case v brs -- if cannot reduce, return neutral term
  eval (LetPair e (Branch b)) = 
    case ppatternMatch (Pat.getPat b) (eval e) of
      Just r -> eval (Pat.instantiate b r)
      Nothing -> error "No match!"
  
  -- | small-step evaluation
  -- >>> step (t1 `App` t0)
  -- Just (λ. λ. 0 (λ. 0 0))
  step :: Exp n -> Maybe (Exp n)
  step (Var x) = Nothing
  step (Lam b) = Nothing
  step (App (Lam b) e2) = Just (instantiate1 b e2)
  step (App e1 e2)
    | Just e1' <- step e1 = Just (App e1' e2)
    | Just e2' <- step e2 = Just (App e1 e2')
    | otherwise = Nothing
  step (LetPair a (Branch b)) 
    | Just r <- ppatternMatch (Pat.getPat b) a
    = Just (Pat.instantiate b r)
  step (LetPair e b) 
    | Just e' <- step e = Just (LetPair e' b)
    | otherwise = Nothing
  step (Con s) = Nothing
  step (Case e brs)
    | Just br <- findBranch e brs = Just br
    | Just e' <- step e = Just (Case e' brs)
    | otherwise = Nothing
  
  eval' :: Exp n -> Exp n
  eval' e
    | Just e' <- step e = eval' e'
    | otherwise = e
  
  -- full normalization
  -- to normalize under a lambda expression, we must first unbind
  -- it and then rebind it when finished
  
  -- >>> nf t1
  -- λ. λ. 1 0
  -- >>> nf (t1 `App` t0)
  -- λ. λ. 0 0
  nf :: Exp n -> Exp n
  nf (Var x) = Var x
  nf (Lam b) = Lam (bind1 (nf (getBody1 b)))
  nf (App e1 e2) =
    case nf e1 of
      Lam b -> instantiate1 b (nf e2)
      t -> App t (nf e2)
  nf (Con s) = Con s
  nf (Case e brs) =
    let v = nf e
     in case findBranch v brs of
          Just b -> nf b
          Nothing -> Case e (map nfBr brs)
  nf (LetPair e br@(Branch b)) = 
    let v = nf e in
    case ppatternMatch (Pat.getPat b) v of 
      Just r -> nf (Pat.instantiate b r)
      Nothing -> LetPair (nf e) (nfBr br)
  
  nfBr :: (forall n. Sized (pat n)) => Branch pat n -> Branch pat n
  nfBr (Branch bnd) =
    Branch (Pat.bind (Pat.getPat bnd) (nf (Pat.getBody bnd)))
  