quox/lib/Quox/FreeVars.idr

310 lines
9 KiB
Idris

module Quox.FreeVars
import Quox.Syntax.Term.Base
import Data.Maybe
import Data.Nat
import Data.Singleton
import Data.SortedSet
import Derive.Prelude
%language ElabReflection
public export
FreeVars' : Nat -> Type
FreeVars' n = Context' Bool n
public export
record FreeVars n where
constructor FV
vars : FreeVars' n
%name FreeVars xs
%runElab deriveIndexed "FreeVars" [Eq, Ord, Show]
export %inline
(||) : FreeVars n -> FreeVars n -> FreeVars n
FV s || FV t = FV $ zipWith (\x, y => x || y) s t
export %inline
(&&) : FreeVars n -> FreeVars n -> FreeVars n
FV s && FV t = FV $ zipWith (\x, y => x && y) s t
export %inline Semigroup (FreeVars n) where (<+>) = (||)
export %inline [AndS] Semigroup (FreeVars n) where (<+>) = (&&)
export
only : {n : Nat} -> Var n -> FreeVars n
only i = FV $ only' i where
only' : {n' : Nat} -> Var n' -> FreeVars' n'
only' VZ = replicate (pred n') False :< True
only' (VS i) = only' i :< False
export %inline
all : {n : Nat} -> FreeVars n
all = FV $ replicate n True
export %inline
none : {n : Nat} -> FreeVars n
none = FV $ replicate n False
export %inline
uncons : FreeVars (S n) -> (FreeVars n, Bool)
uncons (FV (xs :< x)) = (FV xs, x)
export %inline {n : Nat} -> Monoid (FreeVars n) where neutral = none
export %inline [AndM] {n : Nat} -> Monoid (FreeVars n) where neutral = all
private
self : {n : Nat} -> Context' (FreeVars n) n
self = tabulate (\i => FV $ tabulate (== i) n) n
export
shift : forall from, to. Shift from to -> FreeVars from -> FreeVars to
shift by (FV xs) = FV $ shift' by xs where
shift' : Shift from' to' -> FreeVars' from' -> FreeVars' to'
shift' SZ ctx = ctx
shift' (SS by) ctx = shift' by ctx :< False
export
fromSet : {n : Nat} -> SortedSet (Var n) -> FreeVars n
fromSet vs = FV $ tabulateLT n $ \i => contains (V i) vs
export
toSet : {n : Nat} -> FreeVars n -> SortedSet (Var n)
toSet (FV vs) =
foldl_ (\s, i => maybe s (\i => insert i s) i) empty $
zipWith (\i, b => guard b $> i) (tabulateLT n V) vs
public export
interface HasFreeVars (0 tm : Nat -> Type) where
constructor HFV
fv : {n : Nat} -> tm n -> FreeVars n
public export
interface HasFreeDVars (0 tm : TermLike) where
constructor HFDV
fdv : {d, n : Nat} -> tm d n -> FreeVars d
public export %inline
fvWith : HasFreeVars tm => Singleton n -> tm n -> FreeVars n
fvWith (Val n) = fv
public export %inline
fdvWith : HasFreeDVars tm => Singleton d -> Singleton n -> tm d n -> FreeVars d
fdvWith (Val d) (Val n) = fdv
export
Fdv : (0 tm : TermLike) -> {n : Nat} ->
HasFreeDVars tm => HasFreeVars (\d => tm d n)
Fdv tm @{HFDV fdv} = HFV fdv
export
fvEach : {n1, n2 : Nat} -> HasFreeVars env =>
Subst env n1 n2 -> Context' (Lazy (FreeVars n2)) n1
fvEach (Shift by) = map (delay . shift by) self
fvEach (t ::: th) = fvEach th :< fv t
export
fdvEach : {d, n1, n2 : Nat} -> HasFreeDVars env =>
Subst (env d) n1 n2 -> Context' (Lazy (FreeVars d)) n1
fdvEach (Shift by) = replicate n1 none
fdvEach (t ::: th) = fdvEach th :< fdv t
export
HasFreeVars Dim where
fv (K _ _) = none
fv (B i _) = only i
export
{s : Nat} -> HasFreeVars tm => HasFreeVars (Scoped s tm) where
fv (S _ (Y body)) = FV $ drop s (fv body).vars
fv (S _ (N body)) = fv body
export
implementation [DScope]
{s : Nat} -> HasFreeDVars tm =>
HasFreeDVars (\d, n => Scoped s (\d' => tm d' n) d)
where
fdv (S _ (Y body)) = FV $ drop s (fdv body).vars
fdv (S _ (N body)) = fdv body
export
fvD : {0 tm : TermLike} -> {n : Nat} -> (forall d. HasFreeVars (tm d)) =>
Scoped s (\d => tm d n) d -> FreeVars n
fvD (S _ (Y body)) = fv body
fvD (S _ (N body)) = fv body
export
fdvT : HasFreeDVars tm => {s, d, n : Nat} -> Scoped s (tm d) n -> FreeVars d
fdvT (S _ (Y body)) = fdv body
fdvT (S _ (N body)) = fdv body
private
guardM : Monoid a => Bool -> Lazy a -> a
guardM b x = if b then x else neutral
export
implementation
(HasFreeVars tm, HasFreeVars env) =>
HasFreeVars (WithSubst tm env)
where
fv (Sub term subst) =
let Val from = getFrom subst in
foldMap (uncurry guardM) $ zipWith (,) (fv term).vars (fvEach subst)
export
implementation [WithSubst]
((forall d. HasFreeVars (tm d)), HasFreeDVars tm, HasFreeDVars env) =>
HasFreeDVars (\d => WithSubst (tm d) (env d))
where
fdv (Sub term subst) =
let Val from = getFrom subst in
fdv term <+>
foldMap (uncurry guardM) (zipWith (,) (fv term).vars (fdvEach subst))
export HasFreeVars (Term d)
export HasFreeVars (Elim d)
export
HasFreeVars (Term d) where
fv (TYPE {}) = none
fv (IOState {}) = none
fv (Pi {arg, res, _}) = fv arg <+> fv res
fv (Lam {body, _}) = fv body
fv (Sig {fst, snd, _}) = fv fst <+> fv snd
fv (Pair {fst, snd, _}) = fv fst <+> fv snd
fv (Enum {}) = none
fv (Tag {}) = none
fv (Eq {ty, l, r, _}) = fvD ty <+> fv l <+> fv r
fv (DLam {body, _}) = fvD body
fv (NAT {}) = none
fv (Nat {}) = none
fv (Succ {p, _}) = fv p
fv (STRING {}) = none
fv (Str {}) = none
fv (BOX {ty, _}) = fv ty
fv (Box {val, _}) = fv val
fv (Let {rhs, body, _}) = fv rhs <+> fv body
fv (E e) = fv e
fv (CloT s) = fv s
fv (DCloT s) = fv s.term
export
HasFreeVars (Elim d) where
fv (F {}) = none
fv (B i _) = only i
fv (App {fun, arg, _}) = fv fun <+> fv arg
fv (CasePair {pair, ret, body, _}) = fv pair <+> fv ret <+> fv body
fv (Fst pair _) = fv pair
fv (Snd pair _) = fv pair
fv (CaseEnum {tag, ret, arms, _}) =
fv tag <+> fv ret <+> foldMap fv (values arms)
fv (CaseNat {nat, ret, zero, succ, _}) =
fv nat <+> fv ret <+> fv zero <+> fv succ
fv (CaseBox {box, ret, body, _}) =
fv box <+> fv ret <+> fv body
fv (DApp {fun, _}) = fv fun
fv (Ann {tm, ty, _}) = fv tm <+> fv ty
fv (Coe {ty, val, _}) = fvD ty <+> fv val
fv (Comp {ty, val, zero, one, _}) =
fv ty <+> fv val <+> fvD zero <+> fvD one
fv (TypeCase {ty, ret, arms, def, _}) =
fv ty <+> fv ret <+> fv def <+> foldMap (\x => fv x.snd) (toList arms)
fv (CloE s) = fv s
fv (DCloE s) = fv s.term
private
expandDShift : {d1 : Nat} -> Shift d1 d2 -> Loc -> Context' (Dim d2) d1
expandDShift by loc = tabulateLT d1 (\i => BV i loc // by)
private
expandDSubst : {d1 : Nat} -> DSubst d1 d2 -> Loc -> Context' (Dim d2) d1
expandDSubst (Shift by) loc = expandDShift by loc
expandDSubst (t ::: th) loc = expandDSubst th loc :< t
private
fdvSubst' : {d1, d2, n : Nat} -> (Located2 tm, HasFreeDVars tm) =>
tm d1 n -> DSubst d1 d2 -> FreeVars d2
fdvSubst' t th =
fold $ zipWith maybeOnly (fdv t).vars (expandDSubst th t.loc)
where
maybeOnly : {d : Nat} -> Bool -> Dim d -> FreeVars d
maybeOnly True (B i _) = only i
maybeOnly _ _ = none
private
fdvSubst : {d, n : Nat} -> (Located2 tm, HasFreeDVars tm) =>
WithSubst (\d => tm d n) Dim d -> FreeVars d
fdvSubst (Sub t th) = let Val from = getFrom th in fdvSubst' t th
export HasFreeDVars Term
export HasFreeDVars Elim
export
HasFreeDVars Term where
fdv (TYPE {}) = none
fdv (IOState {}) = none
fdv (Pi {arg, res, _}) = fdv arg <+> fdvT res
fdv (Lam {body, _}) = fdvT body
fdv (Sig {fst, snd, _}) = fdv fst <+> fdvT snd
fdv (Pair {fst, snd, _}) = fdv fst <+> fdv snd
fdv (Enum {}) = none
fdv (Tag {}) = none
fdv (Eq {ty, l, r, _}) = fdv @{DScope} ty <+> fdv l <+> fdv r
fdv (DLam {body, _}) = fdv @{DScope} body
fdv (NAT {}) = none
fdv (Nat {}) = none
fdv (Succ {p, _}) = fdv p
fdv (STRING {}) = none
fdv (Str {}) = none
fdv (BOX {ty, _}) = fdv ty
fdv (Box {val, _}) = fdv val
fdv (Let {rhs, body, _}) = fdv rhs <+> fdvT body
fdv (E e) = fdv e
fdv (CloT s) = fdv s @{WithSubst}
fdv (DCloT s) = fdvSubst s
export
HasFreeDVars Elim where
fdv (F {}) = none
fdv (B {}) = none
fdv (App {fun, arg, _}) = fdv fun <+> fdv arg
fdv (CasePair {pair, ret, body, _}) = fdv pair <+> fdvT ret <+> fdvT body
fdv (Fst pair _) = fdv pair
fdv (Snd pair _) = fdv pair
fdv (CaseEnum {tag, ret, arms, _}) =
fdv tag <+> fdvT ret <+> foldMap fdv (values arms)
fdv (CaseNat {nat, ret, zero, succ, _}) =
fdv nat <+> fdvT ret <+> fdv zero <+> fdvT succ
fdv (CaseBox {box, ret, body, _}) =
fdv box <+> fdvT ret <+> fdvT body
fdv (DApp {fun, arg, _}) =
fdv fun <+> fv arg
fdv (Ann {tm, ty, _}) =
fdv tm <+> fdv ty
fdv (Coe {ty, p, q, val, _}) =
fdv @{DScope} ty <+> fv p <+> fv q <+> fdv val
fdv (Comp {ty, p, q, val, r, zero, one, _}) =
fdv ty <+> fv p <+> fv q <+> fdv val <+>
fv r <+> fdv @{DScope} zero <+> fdv @{DScope} one
fdv (TypeCase {ty, ret, arms, def, _}) =
fdv ty <+> fdv ret <+> fdv def <+> foldMap (\x => fdvT x.snd) (toList arms)
fdv (CloE s) = fdv s @{WithSubst}
fdv (DCloE s) = fdvSubst s