quox/lib/Quox/Syntax/Qty.idr

||| quantities count how many times a bound variable is used [@nuttin; @qtt].
|||
||| i tried grtt [@grtt] for a bit but i think it was more complex than
||| it's worth in a language that has other stuff going on too
module Quox.Syntax.Qty

import public Quox.Polynomial
import Quox.Name
import Quox.Syntax.Subst
import Quox.Pretty
import Quox.Decidable
import Quox.PrettyValExtra
import Data.DPair
import Data.Singleton
import Derive.Prelude

%default total
%language ElabReflection


||| the possibilities we care about are:
|||
||| - 0: a variable is used only at compile time, not run time
||| - 1: a variable is used exactly once at run time
||| - ω (or #): don't care. an ω variable *can* also be used 0/1 time
public export
data QConst = Zero | One | Any
%runElab derive "QConst" [Eq, Ord, Show, PrettyVal]
%name QConst q, r


export
prettyQConst : {opts : _} -> QConst -> Eff Pretty (Doc opts)
prettyQConst Zero = hl Qty $ text "0"
prettyQConst One  = hl Qty $ text "1"
prettyQConst Any  = hl Qty =<< ifUnicode (text "ω") (text "#")

-- ||| prints in a form that can be a suffix of "case"
-- public export
-- prettySuffix : {opts : _} -> Qty -> Eff Pretty (Doc opts)
-- prettySuffix = prettyQty


namespace QConst
  ||| e.g. if in the expression `(s, t)`, the variable `x` is
  ||| used π times in `s` and ρ times in `t`, then it's used
  ||| (π + ρ) times in the whole expression
  public export
  plus : QConst -> QConst -> QConst
  plus Zero rh   = rh
  plus pi   Zero = pi
  plus _    _    = Any

  ||| e.g. if a function `f` uses its argument π times,
  ||| and `f x` occurs in a σ context, then `x` is used `πσ` times overall
  public export
  times : QConst -> QConst -> QConst
  times Zero _    = Zero
  times _    Zero = Zero
  times One  rh   = rh
  times pi   One  = pi
  times Any  Any  = Any

  -- ||| "π ∨ ρ"
  -- |||
  -- ||| returns a quantity τ with π ≤ τ and ρ ≤ τ.
  -- ||| if π = ρ, then it's that, otherwise it's ω.
  -- public export
  -- lub : QConst -> QConst -> QConst
  -- lub p q = if p == q then p else Any

export %inline
AddMonoid QConst where zero = Zero; (+.) = plus; isZero = (== Zero)

export %inline
TimesMonoid QConst where one = One; (*.) = times


public export
record Qty n where
  constructor Q
  value : Polynomial QConst n
  loc   : Loc
%runElab deriveIndexed "Qty" [Eq, Ord]

export %hint
ShowQty : {n : Nat} -> Show (Qty n)
ShowQty = deriveShow

export %hint
PrettyValQty : {n : Nat} -> PrettyVal (Qty n)
PrettyValQty = derivePrettyVal


export
zero, one, any : {n : Nat} -> Loc -> Qty n
zero = Q zero
one  = Q one
any  = Q $ pconst Any

export
(+) : Qty n -> Qty n -> Qty n
Q xs l1 + Q ys l2 = Q (xs +. ys) (l1 `or` l2)

export
(*) : Qty n -> Qty n -> Qty n
Q xs l1 * Q ys l2 = Q (xs *. ys) (l1 `or` l2)

export
isAny : Qty n -> Bool
isAny (Q pi _) = pi.at0 == Any

export
lub : {n : Nat} -> Qty n -> Qty n -> Qty n
lub pi rh = if pi == rh then pi else any pi.loc

||| "π ≤ ρ"
|||
||| if a variable is bound with quantity ρ, then it can be used with an actual
||| quantity π as long as π ≤ ρ. for example, an ω variable can be used any
||| number of times, so π ≤ ω for any π.
public export
compat : Qty n -> Qty n -> Bool
compat pi rh = isAny rh || pi == rh


private
toVarString : BContext n -> Monom n -> List BindName
toVarString ns ds = fold $ zipWith replicate ds ns

private
prettyTerm : {opts : LayoutOpts} ->
             BContext n -> Monom n -> QConst -> Eff Pretty (Doc opts)
prettyTerm ns ds pi = do
  pi <- prettyQConst pi
  xs <- traverse prettyQBind (toVarString ns ds)
  pure $ separateTight !qtimesD $ pi :: xs

export
prettyQty : {opts : LayoutOpts} -> BContext n -> Qty n -> Eff Pretty (Doc opts)
prettyQty ns (Q q _) =
  let Val _ = lengthPrf0 ns in
  separateLoose !qplusD <$>
  traverse (uncurry $ prettyTerm ns) (toList q)

||| to maintain subject reduction, only 0 or 1 can occur
||| for the subject of a typing judgment. see @qtt, §2.3 for more detail
public export
data SQty = SZero | SOne
%runElab derive "SQty" [Eq, Ord, Show, PrettyVal]
%name Qty.SQty sg

||| "σ ⨴ π"
|||
||| σ ⨴ π is 0 if either of σ or π are, otherwise it is σ.
public export
subjMult : SQty -> Qty n -> SQty
subjMult sg pi = if isZero pi.value then SZero else sg


||| it doesn't make much sense for a top level declaration to have a
||| quantity of 1, so the only distinction is whether it is present
||| at runtime at all or not
public export
data GQty = GZero | GAny
%runElab derive "GQty" [Eq, Ord, Show, PrettyVal]
%name GQty rh

||| when checking a definition, a 0 definition is checked at 0,
||| but an ω definition is checked at 1 since ω isn't a subject quantity
public export %inline
globalToSubj : GQty -> SQty
globalToSubj GZero = SZero
globalToSubj GAny  = SOne


public export
DecEq SQty where
  decEq SZero SZero = Yes Refl
  decEq SZero SOne  = No $ \case _ impossible
  decEq SOne  SZero = No $ \case _ impossible
  decEq SOne  SOne  = Yes Refl

public export
DecEq GQty where
  decEq GZero GZero = Yes Refl
  decEq GZero GAny  = No $ \case _ impossible
  decEq GAny  GZero = No $ \case _ impossible
  decEq GAny  GAny  = Yes Refl


namespace SQty
  public export %inline
  toQty : {n : Nat} -> Loc -> SQty -> Qty n
  toQty loc SZero = zero loc
  toQty loc SOne  = one  loc

  public export %inline
  (.qconst) : SQty -> QConst
  (SZero).qconst = Zero
  (SOne).qconst  = One

namespace GQty
  public export %inline
  toQty : {n : Nat} -> Loc -> GQty -> Qty n
  toQty loc GZero = zero loc
  toQty loc GAny  = any  loc

  public export %inline
  (.qconst) : GQty -> QConst
  (GZero).qconst = Zero
  (GAny).qconst  = Any


export
prettySQty : {opts : _} -> SQty -> Eff Pretty (Doc opts)
prettySQty sg = prettyQConst sg.qconst

export
prettyGQty : {opts : _} -> GQty -> Eff Pretty (Doc opts)
prettyGQty pi = prettyQConst pi.qconst


public export
QSubst : Nat -> Nat -> Type
QSubst = Subst Qty

export
FromVarR Qty where
  fromVarR i loc = Q (fromVarR i loc) loc

export
CanShift Qty where
  Q p loc // by = Q (p // by) loc

export
CanSubstSelfR Qty where
  Q q loc //? th = Q (q //? map (.value) th) loc