add polynomials

lib/Quox/Polynomial.idr

@@ -0,0 +1,189 @@
+module Quox.Polynomial
+
+import public Quox.Syntax.Subst
+import public Quox.Context
+import public Quox.Semiring
+import Data.Maybe
+import Data.SortedMap
+import Data.Singleton
+import Quox.PrettyValExtra
+
+%default total
+
+%hide Prelude.toList
+
+public export
+Monom : Nat -> Type
+Monom = Context' Nat
+
+export %inline
+mZero : {n : Nat} -> Monom n
+mZero = replicate n 0
+
+export %inline
+mIsZero : Monom n -> Bool
+mIsZero = all (== 0)
+
+
+private
+PolyBody : Type -> Nat -> Type
+PolyBody coeff scope = SortedMap (Monom scope) coeff
+
+private %inline
+getScope : PolyBody coeff scope -> Maybe (Singleton scope)
+getScope p = lengthPrf0 . fst <$> leftMost p
+
+export
+data Polynomial : (coeff : Type) -> (scope : Nat) -> Type where
+  Const : (k : coeff) -> Polynomial coeff scope
+  Many  : (p : PolyBody coeff scope) -> Polynomial coeff scope
+%name Polynomial p, q
+-- the Const constructor exists so that `pconst` (and by extension, `one`, and
+-- the TimesMonoid instance) doesn't need to know the scope size
+
+export %inline
+toConst' : List (Monom scope, coeff) -> Maybe coeff
+toConst' [(m, k)] = k <$ guard (mIsZero m)
+toConst' _        = Nothing
+
+export %inline
+toConst : PolyBody coeff scope -> Maybe coeff
+toConst = toConst' . toList
+
+export %inline
+toTrimmedList : AddMonoid v => SortedMap k v -> List (k, v)
+toTrimmedList = List.filter (not . isZero . snd) . toList
+
+export %inline
+toList : {scope : Nat} -> AddMonoid coeff =>
+         Polynomial coeff scope -> List (Monom scope, coeff)
+toList (Const k) = [(mZero, k)]
+toList (Many  p) = toTrimmedList p
+
+private %inline
+addBody : AddMonoid coeff =>
+          PolyBody coeff scope -> PolyBody coeff scope -> PolyBody coeff scope
+addBody = mergeWith (+.)
+
+export %inline
+fromList : AddMonoid coeff =>
+           List (Monom scope, coeff) -> Polynomial coeff scope
+fromList xs = case toConst' xs of
+  Just k  => Const k
+  Nothing => Many $ foldl add1 empty xs
+where
+  add1 : PolyBody coeff scope -> (Monom scope, coeff) -> PolyBody coeff scope
+  add1 p (x, k) = p `addBody` singleton x k
+
+
+export %inline
+(AddMonoid coeff, Eq coeff) => Eq (Polynomial coeff scope) where
+  Const x == Const y = x == y
+  Const x == Many  q = maybe False (== x) $ toConst q
+  Many  p == Const y = maybe False (== y) $ toConst p
+  Many  p == Many  q = toTrimmedList p == toTrimmedList q
+
+export %inline
+(AddMonoid coeff, Ord coeff) => Ord (Polynomial coeff scope) where
+  compare (Const x) (Const y) = compare x y
+  compare (Const x) (Many  q) = maybe LT (compare x) $ toConst q
+  compare (Many  p) (Const y) = maybe GT (flip compare y) $ toConst p
+  compare (Many  p) (Many  q) = compare (toTrimmedList p) (toTrimmedList q)
+
+export %inline
+{scope : Nat} -> (AddMonoid coeff, Show coeff) =>
+Show (Polynomial coeff scope) where
+  showPrec d p = showCon d "fromList" $ showArg (toList p)
+
+export %inline
+{scope : Nat} -> (AddMonoid coeff, PrettyVal coeff) =>
+PrettyVal (Polynomial coeff scope) where
+  prettyVal p = Con "fromList" [prettyVal $ toList p]
+
+
+export %inline
+pconst : a -> Polynomial a n
+pconst = Const
+
+export %inline
+(.at) : AddMonoid a => Polynomial a n -> Monom n -> a
+(Const k).at m = if mIsZero m then k else zero
+(Many  p).at m = fromMaybe zero $ lookup m p
+
+export %inline
+(.at0) : AddMonoid a => Polynomial a n -> a
+(Const k).at0 = k
+(Many  p).at0 = fromMaybe zero $ do
+  (m, k) <- leftMost p
+  guard $ mIsZero m
+  pure k
+
+
+private %inline
+addConstMany : AddMonoid coeff =>
+               coeff -> PolyBody coeff scope -> Polynomial coeff scope
+addConstMany k p = case getScope p of
+  Nothing      => Const k
+  Just (Val _) => Many $ p `addBody` singleton mZero k
+
+export %inline
+AddMonoid a => AddMonoid (Polynomial a n) where
+  zero   = Const zero
+
+  isZero (Const k) = isZero k
+  isZero (Many  p) = maybe False isZero $ toConst p
+
+  Const j +. Const k = Const $ j +. k
+  Const j +. Many  q = addConstMany j q
+  Many  p +. Const k = addConstMany k p
+  Many  p +. Many  q = Many $ p `addBody` q
+
+export %inline
+Semiring a => TimesMonoid (Polynomial a n) where
+  one = pconst one
+
+  Const j *. Const k = Const $ j *. k
+  Const j *. Many  q = Many $ map (j *.) q
+  Many  p *. Const k = Many $ map (*. k) p
+  Many  p *. Many  q = fromList $ do
+    (xs, i) <- toList p
+    (ys, j) <- toList q
+    let k = i *. j
+    guard $ not $ isZero k
+    pure (zipWith (+) xs ys, i *. j)
+
+export %inline
+Semiring a => VecSpace a (Polynomial a n) where
+  x *: xs = Const x *. xs
+
+
+private
+shiftMonom : Shift from to -> Monom from -> Monom to
+shiftMonom SZ      m = m
+shiftMonom (SS by) m = shiftMonom by m :< 0
+
+export
+CanShift (Polynomial coeff) where
+  Const k // _  = Const k
+  Many  p // by =
+    let xs = mapFst (shiftMonom by) <$> toList p in
+    Many $ fromList xs
+
+
+export
+TimesMonoid a => FromVarR (Polynomial a) where
+  fromVarR i l {n} =
+    let m = tabulateVar n $ \j => if i == j then 1 else 0 in
+    Many $ singleton m one
+
+
+private
+substMonom : Semiring a => {from, to : Nat} ->
+             Subst (Polynomial a) from to -> Monom from -> a -> Polynomial a to
+substMonom th m k = sproduct {init = pconst k} $
+  zipWith (\i, p => getR th i noLoc ^. p) (allVars from) m
+
+export
+Semiring a => CanSubstSelfR (Polynomial a) where
+  Const k //? _  = Const k
+  Many  p //? th = ssum $ map (uncurry $ substMonom th) $ toList p

lib/Quox/Semiring.idr

@@ -0,0 +1,122 @@
+module Quox.Semiring
+
+import Quox.Context
+import public Quox.NatExtra
+import Data.Bool.Decidable
+import Quox.No
+
+%hide Nat.isZero
+
+export infixl 8 +.
+export infixl 9 *., *:, :*
+export infixr 10 ^.
+
+public export
+interface AddMonoid a where
+  zero   : a
+  isZero : a -> Bool
+  (+.)   : a -> a -> a
+
+public export
+interface TimesMonoid a where
+  one : a
+  (*.) : a -> a -> a
+
+public export
+0 Semiring : Type -> Type
+Semiring a = (AddMonoid a, TimesMonoid a)
+
+public export
+interface (Semiring base, AddMonoid vec) => VecSpace base vec | vec where
+  (*:) : base -> vec -> vec
+  (:*) : vec -> base -> vec
+
+  x  *: xs = xs :* x
+  xs :* x  = x  *: xs
+
+
+namespace Foldable
+  public export %inline %tcinline
+  ssum : AddMonoid a => Foldable t => {default zero init : a} -> t a -> a
+  ssum = foldl (+.) init
+
+  public export %inline %tcinline
+  sproduct : TimesMonoid a => Foldable t => {default one init : a} -> t a -> a
+  sproduct = foldl (*.) init
+
+namespace Context
+  public export %inline %tcinline
+  ssum : AddMonoid a => {default zero init : a} ->
+         Telescope' a from to -> a
+  ssum = foldl_ (+.) zero
+
+  public export %inline %tcinline
+  sproduct : TimesMonoid a => {default one init : a} ->
+             Telescope' a from to -> a
+  sproduct = foldl_ (*.) one
+
+
+public export
+(^.) : TimesMonoid a => a -> Nat -> a
+x ^. 0   = one
+x ^. S k = x *. x^.k
+
+
+export %inline
+[NumAdd] Eq a => Num a => AddMonoid a where
+  zero = 0; (+.) = (+); isZero = (== 0)
+
+export %inline
+[NumTimes] Num a => TimesMonoid a where
+  one = 1; (*.) = (*)
+
+
+export %inline %hint
+NatAM : AddMonoid Nat
+NatAM = NumAdd
+
+export %inline %hint
+NatTM : TimesMonoid Nat
+NatTM = NumTimes
+
+export %inline %hint
+IntegerAM : AddMonoid Integer
+IntegerAM = NumAdd
+
+export %inline %hint
+IntegerTM : TimesMonoid Integer
+IntegerTM = NumTimes
+
+export %inline %hint
+IntAM : AddMonoid Int
+IntAM = NumAdd
+
+export %inline %hint
+IntTM : TimesMonoid Int
+IntTM = NumTimes
+
+
+export %inline
+(from `LTE'` to, AddMonoid a) => AddMonoid (Telescope' a from to) where
+  zero   = replicateLTE zero
+  (+.)   = zipWith (+.)
+  isZero = all isZero
+
+export %inline
+(from `LTE'` to, Semiring a) => VecSpace a (Telescope' a from to) where
+  x *: xs = map (x *.) xs
+
+
+public export
+interface AddMonoid a => MonoAddMonoid a where
+  0 zeroPlus : (x, y : a) ->
+               No (isZero x) -> No (isZero y) -> No (isZero (x +. y))
+
+export
+MonoAddMonoid Nat where
+  zeroPlus (S _) (S _) Ah Ah = Ah
+
+
+public export
+0 MonoAddSemiring : Type -> Type
+MonoAddSemiring a = (MonoAddMonoid a, TimesMonoid a)

lib/quox-lib.ipkg

@@ -21,6 +21,8 @@ modules =
   Quox.No,
   Quox.Log,
   Quox.Loc,
+  Quox.Semiring,
+  Quox.Polynomial,
   Quox.Var,
   Quox.Scoped,
   Quox.OPE,