path: root/src/Obs/NormalForm/Normalise.idr

module Obs.NormalForm.Normalise

import Data.Vect

import Obs.Logging
import Obs.NormalForm
import Obs.Sort
import Obs.Substitution

-- Substitution ----------------------------------------------------------------

mergeName : String -> String -> String
mergeName "" s' = s'
mergeName "_" s' = s'
mergeName s s' = s

wkn : Vect k String -> (Fin n -> Logging ann (NormalForm m)) -> Fin (k + n) -> Logging ann (NormalForm (k + m))
wkn []            f = f
wkn (var :: vars) f =
  add
    (Logging ann . NormalForm)
    (pure $ Ntrl $ Var var FZ)
    (\i => pure $ rename !(wkn vars f i) FS)

partial
substCnstr : Constructor n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (Constructor m)
partial
substNtrl  : Neutral n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
partial export
subst      : NormalForm n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
partial export
subst1     : NormalForm n -> NormalForm (S n) -> Logging ann (NormalForm n)
partial export
doApp      : NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial export
doFst      : NormalForm n -> Logging ann (NormalForm n)
partial export
doSnd      : NormalForm n -> Logging ann (NormalForm n)
partial export
doCase     : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial export
doEqual    : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doEqualL   : Nat -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doEqualR   : Nat -> Constructor n -> NormalForm n -> Logging ann (NormalForm n)
partial
doEqualU   : Nat -> Constructor n -> Constructor n -> Logging ann (NormalForm n)
partial export
doCastL    : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doCastR    : Constructor n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
partial
doCastU    : Constructor n -> Constructor n -> NormalForm n -> Logging ann (NormalForm n)

substCnstr (Sort s) f = pure $ Sort s
substCnstr (Pi s s' var a b) f = do
  a <- subst a f
  b <- subst b (wkn [var] f)
  pure (Pi s s' var a b)
substCnstr (Lambda var t) f = do
  t <- subst t (wkn [var] f)
  pure (Lambda var t)
substCnstr (Sigma s s' var a b) f = do
  a <- subst a f
  b <- subst b (wkn [var] f)
  pure (Sigma s s' var a b)
substCnstr (Pair t u) f = do
  t <- subst t f
  u <- subst u f
  pure (Pair t u)
substCnstr (Sum s s' a b) f = do
  a <- subst a f
  b <- subst b f
  pure (Sum s s' a b)
substCnstr (Left t) f = do
  t <- subst t f
  pure (Left t)
substCnstr (Right t) f = do
  t <- subst t f
  pure (Right t)
substCnstr Top f = pure $ Top
substCnstr Bottom  f = pure $ Bottom

substNtrl (Var var i) f = do
  Ntrl (Var var' j) <- f i
    | t => pure t
  pure (Ntrl (Var (mergeName var' var) j))
substNtrl (App t u) f = do
  t <- substNtrl t f
  u <- subst u f
  doApp t u
substNtrl (Fst t) f = do
  t <- substNtrl t f
  doFst t
substNtrl (Snd t) f = do
  t <- substNtrl t f
  doSnd t
substNtrl (Case t u t') f = do
  t <- substNtrl t f
  u <- subst u f
  t' <- subst t' f
  doCase t u t'
substNtrl Absurd f = pure $ Ntrl Absurd
substNtrl (Equal a t u) f = do
  a <- substNtrl a f
  t <- subst t f
  u <- subst u f
  doEqual a t u
substNtrl (EqualL i t u) f = do
  t <- substNtrl t f
  u <- subst u f
  doEqualL i t u
substNtrl (EqualR i t u) f = do
  t <- substCnstr t f
  u <- substNtrl u f
  doEqualR i t u
substNtrl (EqualU i t u) f = do
  t <- substCnstr t f
  u <- substCnstr u f
  doEqualU i t u
substNtrl (CastL a b t) f = do
  a <- substNtrl a f
  b <- subst b f
  t <- subst t f
  doCastL a b t
substNtrl (CastR a b t) f = do
  a <- substCnstr a f
  b <- substNtrl b f
  t <- subst t f
  doCastR a b t
substNtrl (CastU a b t) f = do
  a <- substCnstr a f
  b <- substCnstr b f
  t <- subst t f
  doCastU a b t

subst (Ntrl t)  f = substNtrl t f
subst (Cnstr t) f = map Cnstr $ substCnstr t f
subst Irrel     f = pure Irrel

subst1 t u = subst u (add (Logging ann . NormalForm) (pure t) (pure . point))

doApp (Ntrl t) u = pure $ Ntrl (App t u)
doApp Irrel u = pure $ Irrel
doApp (Cnstr (Lambda var t)) u = subst t (add (Logging ann . NormalForm) (pure u) (pure . point))
doApp (Cnstr t) u = inScope "bug" $ inScope "wrong constructor in apply" $ fatal t

doFst (Ntrl t) = pure $ Ntrl (Fst t)
doFst Irrel = pure $ Irrel
doFst (Cnstr (Pair t u)) = pure $ t
doFst (Cnstr t) = inScope "bug" $ inScope "wrong constructor in fst" $ fatal t

doSnd (Ntrl t) = pure $ Ntrl (Snd t)
doSnd Irrel = pure $ Irrel
doSnd (Cnstr (Pair t u)) = pure $ u
doSnd (Cnstr t) = inScope "bug" $ inScope "wrong constructor in snd" $ fatal t

doCase (Ntrl t) f g = pure $ Ntrl (Case t f g)
doCase Irrel f g = inScope "bug" $ inScope "wrong constructor in case" $ fatal "Irrel"
doCase (Cnstr (Left t)) f g = doApp f t
doCase (Cnstr (Right t)) f g = doApp g t
doCase (Cnstr t) f g = inScope "bug" $ inScope "wrong constructor in case" $ fatal t

doEqual (Ntrl a) t u = pure $ Ntrl (Equal a t u)
doEqual Irrel t u = inScope "bug" $ inScope "wrong type over equal" $ fatal "Irrel"
doEqual (Cnstr (Sort Prop)) t u =
  pure $
  Cnstr (Sigma Prop Prop "_"
    (Cnstr $ Pi Prop Prop "_" t (weaken 1 u))
    (weaken 1 $ Cnstr $ Pi Prop Prop "_" u (weaken 1 t)))
doEqual (Cnstr (Sort (Set k))) t u = doEqualL k t u
doEqual (Cnstr (Pi s s' var a b)) t u = do
  eqLhs <- doApp (weaken 1 t) (Ntrl $ Var var FZ)
  eqRhs <- doApp (weaken 1 u) (Ntrl $ Var var FZ)
  eq <- doEqual b eqLhs eqRhs -- b in Set because Pi in Set.
  pure $ Cnstr (Pi s Prop var a eq)
doEqual (Cnstr (Sigma s s' var a b)) t u = do
  t1 <- doFst t
  u1 <- doFst u
  eq1 <- case s of
    Prop => pure $ Cnstr Top
    (Set i) => doEqual a t1 u1
  eq2 <- case s' of
    Prop => pure $ Cnstr Top
    (Set i) => do
      bt <- subst1 t1 b
      bu <- subst1 u1 b
      t2 <- doSnd t
      t2 <- doCastL bt bu t2
      u2 <- doSnd u
      doEqual bu t2 u2
  pure $ Cnstr (Sigma Prop Prop "_" eq1 (weaken 1 eq2))
doEqual (Cnstr (Sum s s' a b)) t u = do
  let x = Ntrl $ Var "x" 1
  let y = Ntrl $ Var "y" 0
  ll <- doEqual (weaken 2 a) x y
  lr <- pure $ Cnstr Bottom
  rl <- pure $ Cnstr Bottom
  rr <- doEqual (weaken 2 b) x y
  fBody <- doCase (weaken 1 u) (Cnstr (Lambda "y" ll)) (Cnstr (Lambda "y" lr))
  gBody <- doCase (weaken 1 u) (Cnstr (Lambda "y" rl)) (Cnstr (Lambda "y" rr))
  doCase t (Cnstr (Lambda "x" fBody)) (Cnstr (Lambda "x" gBody))
doEqual (Cnstr Top) t u = pure $ Cnstr Top
doEqual (Cnstr Bottom) t u = pure $ Cnstr Top

doEqualL i (Ntrl t) u = pure $ Ntrl (EqualL i t u)
doEqualL i Irrel u = inScope "bug" $ inScope "wrong type under equalL" $ fatal "Irrel"
doEqualL i (Cnstr t) u = doEqualR i t u

doEqualR i t (Ntrl u) = pure $ Ntrl (EqualR i t u)
doEqualR i t Irrel = inScope "bug" $ inScope "wrong type under equalR" $ fatal "Irrel"
doEqualR i t (Cnstr u) = doEqualU i t u

doEqualU i (Sort s) (Sort s') = pure $ Cnstr Top -- have suc s = i = suc s', and suc injective
doEqualU i (Pi s s' _ a b) (Pi l l' var a' b') = case (s == s' && l == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a' a
    let x = Ntrl $ Var var FZ
    b <- subst b (add (Logging ann . NormalForm) (doCastL (weaken 1 a') (weaken 1 a) x) (pure . Ntrl . Var "" . FS))
    eqRhs <- doEqual (cast s') b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs $ weaken 1 $ Cnstr (Pi s Prop var a' eqRhs))
doEqualU i (Sigma s s' _ a b) (Sigma l l' var a' b') = case (s == s' && l == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a' a
    let x = Ntrl $ Var var FZ
    b <- subst b (add (Logging ann . NormalForm) (doCastL (weaken 1 a') (weaken 1 a) x) (pure . Ntrl . Var "" . FS))
    eqRhs <- doEqual (cast s') b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs $ weaken 1 $ Cnstr (Pi s Prop var a' eqRhs))
doEqualU i (Sum s s' a b) (Sum l l' a' b') = case (s == s' && l == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a a'
    eqRhs <- doEqual (cast s) b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs (weaken 1 eqRhs))
doEqualU i t u = pure $ Ntrl (EqualU i t u) -- assumption: only logical values reach this far

doCastL (Ntrl a) b t = pure $ Ntrl (CastL a b t)
doCastL Irrel b t = inScope "bug" $ inScope "wrong type for cast" $ fatal "Irrel"
doCastL (Cnstr a) b t = doCastR a b t

doCastR a (Ntrl b) t = pure $ Ntrl (CastR a b t)
doCastR a Irrel t = inScope "bug" $ inScope "wrong type for cast" $ fatal "Irrel"
doCastR a (Cnstr b) t = doCastU a b t

doCastU (Sort s) (Sort s') t = pure t
doCastU (Pi s s' _ a b) (Pi l l' var a' b') t = do
  let x' = Ntrl $ Var var FZ
  let x = doCastL (weaken 1 a') (weaken 1 a) x'
  b <- subst b (add (Logging ann . NormalForm) x (pure . Ntrl . Var "" . FS))
  b' <- subst b' (add (Logging ann . NormalForm) (pure x') (pure . Ntrl . Var "" . FS))
  fx <- doApp (weaken 1 t) !x
  cast <- doCastL b b' fx
  pure $ Cnstr (Lambda var cast)
doCastU (Sigma s s' _ a b) (Sigma l l' var a' b') t = do
  t1 <- doFst t
  t2 <- doSnd t
  t1' <- doCastL a a' t1
  b <- subst1 t1 b
  b' <- subst1 t1' b'
  t2' <- doCastL b b' t2
  pure $ Cnstr (Pair t1' t2')
doCastU (Sum s s' a b) (Sum l l' a' b') t = do
  let x = Ntrl $ Var "x" 0
  castL <- doCastL (weaken 1 a) (weaken 1 a') x
  castR <- doCastL (weaken 1 b) (weaken 1 b') x
  doCase t (Cnstr (Lambda "x" (Cnstr (Left castL)))) (Cnstr (Lambda "x" (Cnstr (Right castR))))
doCastU Top Top t = pure Irrel
doCastU Bottom Bottom t = pure Irrel
doCastU a b t = pure $ Ntrl (CastU a b t)