1 files changed, 30 insertions, 0 deletions
diff --git a/src/Language.hs b/src/Language.hs
index d2eee7f..4e6d604 100644
--- a/src/Language.hs
+++ b/src/Language.hs
@@ -17,6 +17,7 @@ module Language (
 
 import Array
 import AST
+import AST.Sparse.Types
 import AST.Types
 import CHAD.Types
 import Data
@@ -69,6 +70,15 @@ inr = NEInr knownTy
 case_ :: NExpr env (TEither a b) -> (Var name1 a :-> NExpr ('(name1, a) : env) c) -> (Var name2 b :-> NExpr ('(name2, b) : env) c) -> NExpr env c
 case_ e (v1 :-> e1) (v2 :-> e2) = NECase e v1 e1 v2 e2
 
+nothing :: KnownTy a => NExpr env (TMaybe a)
+nothing = NENothing knownTy
+
+just :: NExpr env a -> NExpr env (TMaybe a)
+just = NEJust
+
+maybe_ :: NExpr env b -> (Var name a :-> NExpr ('(name, a) : env) b) -> NExpr env (TMaybe a) -> NExpr env b
+maybe_ a (v :-> b) c = NEMaybe a v b c
+
 constArr_ :: forall t n env. (KnownNat n, KnownScalTy t) => Array n (ScalRep t) -> NExpr env (TArr n (TScal t))
 constArr_ x =
   let ty = knownScalTy
@@ -140,6 +150,9 @@ infixl 9 !
 shape :: NExpr env (TArr n t) -> NExpr env (Tup (Replicate n TIx))
 shape = NEShape
 
+length_ :: NExpr env (TArr N1 t) -> NExpr env TIx
+length_ e = snd_ (shape e)
+
 oper :: SOp a t -> NExpr env a -> NExpr env t
 oper = NEOp
 
@@ -157,6 +170,19 @@ custom :: (Var n1 a :-> Var n2 b :-> NExpr ['(n2, b), '(n1, a)] t)
 custom (n1 :-> n2 :-> a) (nf1 :-> nf2 :-> b) (nr1 :-> nr2 :-> c) e1 e2 =
   NECustom n1 n2 a nf1 nf2 b nr1 nr2 c e1 e2
 
+recompute :: NExpr env a -> NExpr env a
+recompute = NERecompute
+
+with :: forall t a env acname. KnownMTy t => NExpr env t -> (Var acname (TAccum t) :-> NExpr ('(acname, TAccum t) : env) a) -> NExpr env (TPair a t)
+with a (n :-> b) = NEWith (knownMTy @t) a n b
+
+accum :: KnownMTy t => SAcPrj p t a -> NExpr env (AcIdxD p t) -> NExpr env a -> NExpr env (TAccum t) -> NExpr env TNil
+accum p a b c = NEAccum knownMTy p a (spDense (acPrjTy p knownMTy)) b c
+
+accumS :: KnownMTy t => SAcPrj p t a -> NExpr env (AcIdxD p t) -> Sparse a b -> NExpr env b -> NExpr env (TAccum t) -> NExpr env TNil
+accumS p a sp b c = NEAccum knownMTy p a sp b c
+
+
 (.==) :: (KnownScalTy st, ScalIsNumeric st ~ True) => NExpr env (TScal st) -> NExpr env (TScal st) -> NExpr env (TScal TBool)
 a .== b = oper (OEq knownScalTy) (pair a b)
 infix 4 .==
@@ -188,6 +214,10 @@ or_ :: NExpr env (TScal TBool) -> NExpr env (TScal TBool) -> NExpr env (TScal TB
 or_ = oper2 OOr
 infixr 2 `or_`
 
+mod_ :: (ScalIsIntegral a ~ True, KnownScalTy a) => NExpr env (TScal a) -> NExpr env (TScal a) -> NExpr env (TScal a)
+mod_ = oper2 (OMod knownScalTy)
+infixl 7 `mod_`
+
 -- | The first alternative is the True case; the second is the False case.
 if_ :: NExpr env (TScal TBool) -> NExpr env t -> NExpr env t -> NExpr env t
 if_ e a b = case_ (oper OIf e) (#_ :-> NEDrop SZ a) (#_ :-> NEDrop SZ b)