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