quox/lib/Quox/Syntax/DimEq.idr

279 lines
7 KiB
Idris

module Quox.Syntax.DimEq
import public Quox.Syntax.Var
import public Quox.Syntax.Dim
import public Quox.Syntax.Subst
import public Quox.Context
import Quox.Pretty
import Quox.Name
import Data.Maybe
import Data.Nat
import Data.DPair
import Data.Fun.Graph
import Decidable.Decidable
import Decidable.Equality
import Derive.Prelude
%language ElabReflection
%default total
public export
DimEq' : Nat -> Type
DimEq' = Context (Maybe . Dim)
public export
data DimEq : Nat -> Type where
ZeroIsOne : DimEq d
C : (eqs : DimEq' d) -> DimEq d
%name DimEq eqs
%runElab deriveIndexed "DimEq" [Eq, Ord, Show]
public export
consistent : DimEq d -> Bool
consistent ZeroIsOne = False
consistent (C eqs) = True
public export
data IfConsistent : DimEq d -> Type -> Type where
Nothing : IfConsistent ZeroIsOne a
Just : a -> IfConsistent (C eqs) a
export
Functor (IfConsistent eqs) where
map f Nothing = Nothing
map f (Just x) = Just (f x)
export
Foldable (IfConsistent eqs) where
foldr f z Nothing = z
foldr f z (Just x) = f x z
export
Traversable (IfConsistent eqs) where
traverse f Nothing = pure Nothing
traverse f (Just x) = Just <$> f x
public export
ifConsistent : Applicative f => (eqs : DimEq d) -> f a -> f (IfConsistent eqs a)
ifConsistent ZeroIsOne act = pure Nothing
ifConsistent (C _) act = Just <$> act
public export
toMaybe : IfConsistent eqs a -> Maybe a
toMaybe Nothing = Nothing
toMaybe (Just x) = Just x
export
fromGround' : Context' DimConst d -> DimEq' d
fromGround' [<] = [<]
fromGround' (ctx :< e) = fromGround' ctx :< Just (K e)
export
fromGround : Context' DimConst d -> DimEq d
fromGround = C . fromGround'
public export %inline
zeroEq : DimEq 0
zeroEq = C [<]
public export %inline
new' : {d : Nat} -> DimEq' d
new' {d = 0} = [<]
new' {d = S d} = new' :< Nothing
public export %inline
new : {d : Nat} -> DimEq d
new = C new'
public export %inline
get' : DimEq' d -> Var d -> Maybe (Dim d)
get' = getWith $ \p, by => map (// by) p
public export %inline
getVar : DimEq' d -> Var d -> Dim d
getVar eqs i = fromMaybe (B i) $ get' eqs i
public export %inline
getShift' : Shift len out -> DimEq' len -> Var len -> Maybe (Dim out)
getShift' = getShiftWith $ \p, by => map (// by) p
public export %inline
get : DimEq' d -> Dim d -> Dim d
get _ (K e) = K e
get eqs (B i) = getVar eqs i
public export %inline
equal : DimEq d -> (p, q : Dim d) -> Bool
equal ZeroIsOne p q = True
equal (C eqs) p q = get eqs p == get eqs q
infixl 7 :<?
export %inline
(:<?) : DimEq d -> Maybe (Dim d) -> DimEq (S d)
ZeroIsOne :<? d = ZeroIsOne
C eqs :<? d = C $ eqs :< map (get eqs) d
private %inline
ifVar : Var d -> Dim d -> Maybe (Dim d) -> Maybe (Dim d)
ifVar i p = map $ \q => if q == B i then p else q
-- (using decEq instead of (==) because of the proofs below)
private %inline
checkConst : (e, f : DimConst) -> (eqs : Lazy (DimEq' d)) -> DimEq d
checkConst e f eqs = if isYes $ e `decEq` f then C eqs else ZeroIsOne
export
setConst : Var d -> DimConst -> DimEq' d -> DimEq d
setConst VZ e (eqs :< Nothing) = C $ eqs :< Just (K e)
setConst VZ e (eqs :< Just (K f)) = checkConst e f $ eqs :< Just (K f)
setConst VZ e (eqs :< Just (B i)) = setConst i e eqs :<? Just (K e)
setConst (VS i) e (eqs :< p) = setConst i e eqs :<? ifVar i (K e) p
mutual
private
setVar' : (i, j : Var d) -> (0 _ : i `LT` j) -> DimEq' d -> DimEq d
setVar' VZ (VS i) LTZ (eqs :< Nothing) =
C eqs :<? Just (B i)
setVar' VZ (VS i) LTZ (eqs :< Just (K e)) =
setConst i e eqs :<? Just (K e)
setVar' VZ (VS i) LTZ (eqs :< Just (B j)) =
setVar i j eqs :<? Just (B (max i j))
setVar' (VS i) (VS j) (LTS lt) (eqs :< p) =
setVar' i j lt eqs :<? ifVar i (B j) p
export %inline
setVar : (i, j : Var d) -> DimEq' d -> DimEq d
setVar i j eqs with (compareP i j) | (compare i.nat j.nat)
setVar i j eqs | IsLT lt | LT = setVar' i j lt eqs
setVar i i eqs | IsEQ | EQ = C eqs
setVar i j eqs | IsGT gt | GT = setVar' j i gt eqs
export %inline
set : (p, q : Dim d) -> DimEq d -> DimEq d
set _ _ ZeroIsOne = ZeroIsOne
set (K e) (K f) (C eqs) = checkConst e f eqs
set (K e) (B i) (C eqs) = setConst i e eqs
set (B i) (K e) (C eqs) = setConst i e eqs
set (B i) (B j) (C eqs) = setVar i j eqs
public export %inline
Split : Nat -> Type
Split d = (DimEq' d, DSubst (S d) d)
export %inline
split1 : DimConst -> DimEq' (S d) -> Maybe (Split d)
split1 e eqs = case setConst VZ e eqs of
ZeroIsOne => Nothing
C (eqs :< _) => Just (eqs, K e ::: id)
export %inline
split : DimEq' (S d) -> List (Split d)
split eqs = toList (split1 Zero eqs) <+> toList (split1 One eqs)
export
splits' : DimEq' d -> List (DSubst d 0)
splits' [<] = [id]
splits' eqs@(_ :< _) = [th . ph | (eqs', th) <- split eqs, ph <- splits' eqs']
export %inline
splits : DimEq d -> List (DSubst d 0)
splits ZeroIsOne = []
splits (C eqs) = splits' eqs
private
0 newGetShift : (d : Nat) -> (i : Var d) -> (by : Shift d d') ->
getShift' by (new' {d}) i = Nothing
newGetShift (S d) VZ by = Refl
newGetShift (S d) (VS i) by = newGetShift d i (ssDown by)
export
0 newGet' : (d : Nat) -> (i : Var d) -> get' (new' {d}) i = Nothing
newGet' d i = newGetShift d i SZ
export
0 newGet : (d : Nat) -> (p : Dim d) -> get (new' {d}) p = p
newGet d (K e) = Refl
newGet d (B i) = rewrite newGet' d i in Refl
export
0 setSelf : (p : Dim d) -> (eqs : DimEq d) -> set p p eqs = eqs
setSelf p ZeroIsOne = Refl
setSelf (K Zero) (C eqs) = Refl
setSelf (K One) (C eqs) = Refl
setSelf (B i) (C eqs) with (compareP i i) | (compare i.nat i.nat)
_ | IsLT lt | LT = absurd lt
_ | IsEQ | EQ = Refl
_ | IsGT gt | GT = absurd gt
private
prettyDVars : Pretty.HasEnv m => m (SnocList (Doc HL))
prettyDVars = map (pretty0 False . DV) <$> asks dnames
private
prettyCst : (PrettyHL a, PrettyHL b, Pretty.HasEnv m) => a -> b -> m (Doc HL)
prettyCst p q = do
p <- pretty0M p
q <- pretty0M q
pure $ hsep [p, hl Syntax "=", q]
export
PrettyHL (DimEq' d) where
prettyM eqs {m} = do
vars <- prettyDVars
eqs <- go eqs
let prec = if length vars <= 1 && null eqs then Arg else Outer
parensIfM prec $ align $ fillSep $ punctuate comma $ toList $ vars ++ eqs
where
tail : SnocList a -> SnocList a
tail [<] = [<]
tail (xs :< _) = xs
go : DimEq' d' -> m (SnocList (Doc HL))
go [<] = pure [<]
go (eqs :< Nothing) = local {dnames $= tail} $ go eqs
go (eqs :< Just p) = do
eq <- prettyCst (BV {d = 1} 0) (weakD 1 p)
eqs <- local {dnames $= tail} $ go eqs
pure $ eqs :< eq
export
PrettyHL (DimEq d) where
prettyM ZeroIsOne = parensIfM Outer $
align $ fillSep $ punctuate comma $ toList $
!prettyDVars :< !(prettyCst Zero One)
prettyM (C eqs) = prettyM eqs
export
prettyDimEq : NContext d -> DimEq d -> Doc HL
prettyDimEq ds = pretty0With False (toSnocList' ds) [<]
public export
wf' : DimEq' d -> Bool
wf' [<] = True
wf' (eqs :< Nothing) = wf' eqs
wf' (eqs :< Just (K e)) = wf' eqs
wf' (eqs :< Just (B i)) = isNothing (get' eqs i) && wf' eqs
public export
wf : DimEq d -> Bool
wf ZeroIsOne = True
wf (C eqs) = wf' eqs