Algebra.hs

{-# OPTIONS_GHC -Wall -fno-warn-orphans -fno-warn-unused-top-binds -fno-warn-name-shadowing #-}
{-# LANGUAGE FunctionalDependencies, OverloadedStrings #-}

import Data.List hiding (singleton)
import Data.Map
import Data.Semigroup
import Data.String
import GHC.Num.Integer qualified
import GHC.Real qualified
import Numeric.Natural
import Prelude hiding ((^))

-- I'm defining this merely so that the second argument won't be polymorphic.
-- The `(^)` from `Prelude` has signature `a -> b -> a`, but the `b` there is impossible to infer.
(^) :: Num a => a -> Natural -> a
(^) x n = (GHC.Real.^) x (GHC.Num.Integer.integerFromNatural n)

-- If `a` is a `Num` type, then `a`-valued functions is also a `Num` type in a natural way.
instance Num a => Num (x -> a) where
  -- the sum of two functions is the function that plugs x into each term function and adds the two results.
  f + g = \x -> f x + g x
  -- the product of two functions is the function that plugs x into each factor function and multiplies the two results.
  f * g = \x -> f x * g x
  negate f = \x -> negate (f x)
  abs f = \x -> abs (f x)
  signum f = \x -> signum (f x)
  -- an ordinary integer can be thought of as a constant function
  fromInteger n = \_ -> fromInteger n

-- An `r`-algebra is a `Num` type `a` extended by introduction of scalar multiplication by a `Num` type `r`.
-- Common examples are sets of square matrices.
class (Num r, Num a) => Algebra r a | a -> r where
  scale :: r -> a -> a

-- `a`-valued functions is an `a`-algebra when `a` is a `Num` type.
instance Num a => Algebra a (x -> a) where
  scale a f = \x -> a * f x

-- The set of `r`-polynomials in `x`
-- i.e. polynomials with coefficients in `r` and variables in `x`
newtype Polynomial r x
  -- An `r`-polynomial is a function that "makes sense" in any `r`-algebra.
  -- The `(x -> a)` is the dictionary of assignments of `a`-values to the polynomial's variables.
  -- If `x` has only two elements, for example, you would construct an `x -> a` by specifying two `a` values, one for each element of `x`.
  -- In that case, the signature `(x -> a) -> a` would be equivalent to `(a, a) -> a`.
  -- Thus, the signature `(x -> a) -> a` means "an `a`-valued function with `x`-many `a`-valued variables."
  -- The `forall a. Algebra r a =>` part of the signature means that the `(x -> a) -> a` part "works" for any type `a` so long as `a` is an `r`-algebra.
  -- Thus, an `r`-polynomial in `x` is a function with `x`-many variables that "makes sense" to interpret in any `r`-algebra.
  = Polynomial (forall a. Algebra r a => (x -> a) -> a)
  -- Oh, yeah, `Polynomial r` is a functor, as well.
  -- If you have an `r`-polynomial with varibles in `x`, you can convert it so that its variables are in `z` just by replacing the `x` values with `z` values.
  deriving Functor

-- The set of polynomials is always a `Num` type, regardless of choice of `r` or `x`.
instance Num (Polynomial r x) where
  Polynomial p1 + Polynomial p2 = Polynomial (p1 + p2)
  Polynomial p1 * Polynomial p2 = Polynomial (p1 * p2)
  negate (Polynomial p) = Polynomial (negate p)
  abs (Polynomial p) = Polynomial (abs p)
  signum (Polynomial p) = Polynomial (signum p)
  fromInteger n = Polynomial (fromInteger n)

-- If `r` is a `Num` type, then the set of `r`-polynomials in `x` is an `r`-algebra, regardless of choice of `x`.
instance Num r => Algebra r (Polynomial r x) where
  scale r (Polynomial p) = Polynomial (scale r . p)

-- This type is for proving constraints are met.
-- Using `Witness` in your code will only compile if the constraint `c a` is met.
-- The compiler will issue an error if `Witness` is used and the constraint is not met.
data Witness c a where
  Witness :: forall c a. c a => Witness c a

-- The class of free object constructions.
-- A free object construction `f` of a class `c` is a functor whose range consists of free `c` objects.
-- https://en.wikipedia.org/wiki/Free_object
-- https://wiki.haskell.org/Free_structure
-- `f x` is called "the free `c` generated by `x`".
-- The word "free" is like "free of context" or "free of external meaning" and comes from the term "free variable".
class Functor f => Free c f | f -> c where
  -- `f x` must satisfy `c`, regardless of choice of `x`.
  free :: Witness c (f x)
  -- members of `x` are free variables in `f x`
  var :: x -> f x
  -- a member of `f x` can be resolved to a member of `a` by assigning values to the free variables.
  eval :: c a => (x -> a) -> f x -> a

-- The set of `r`-polynomials in `x` is the free `r`-algebra generated by `x`.
instance Num r => Free (Algebra r) (Polynomial r) where
  free = Witness
  var x = Polynomial ($ x)
  eval d (Polynomial p) = p d

-- A standard form for monomials.
-- A monomial is a product of powers of free variables, so we encode them as `Map`s.
-- The `Map` tells us which variables and what their exponents are.
newtype MonomialNormalForm x = MonomialNormalForm (Map x Natural)
  deriving (Eq, Ord)

mnf :: x -> MonomialNormalForm x
mnf x = MonomialNormalForm $ singleton x 1

-- We need a normal form for monomials so that we can display them as text in a standard way.
instance Show x => Show (MonomialNormalForm x) where
  show (MonomialNormalForm m) = foldMap (\(x,n) -> show x <> if n == 1 then "" else show n) $ assocs m

-- The set of monomials is a semigroup.
instance Ord x => Semigroup (MonomialNormalForm x) where
  MonomialNormalForm m1 <> MonomialNormalForm m2 = MonomialNormalForm (unionWith (+) m1 m2)
  stimes = stimesMonoid

-- The set of monomials is a monoid.
instance Ord x => Monoid (MonomialNormalForm x) where
  mempty = MonomialNormalForm mempty

-- A standard form for polynomials.
-- A polynomial is a sum of monomials with `coefficients` in `r`, so we encode them as `Map`s.
-- The `Map` tells us which monomials and what their coefficients are.
newtype PolynomialNormalForm r x = PolynomialNormalForm (Map (MonomialNormalForm x) r)
  deriving (Eq, Ord)

pnf :: Num r => x -> PolynomialNormalForm r x
pnf x = PolynomialNormalForm $ singleton (mnf x) 1

-- The set of polynomials is a `Num` type, and so their normal-form encoding should be, too.
-- The one subtle thing is we need to "normalize" after certain operations in order to ensure that we never have any "duplicate" representations of the same polynomials.
-- This is because two different `Map` values can represent the same polynomial (e.g. insert extra map keys but with values `0`, like adding zero terms to a sum).
-- We need to prevent that if we want `Show` and `Eq` and `Ord` to work correctly, so we normalize by removing zero terms after operations that might introduce them.
instance (Num r, Eq r, Ord x) => Num (PolynomialNormalForm r x) where
  -- This is the "Combine Like Terms" rule, written as a computer program.
  PolynomialNormalForm p1 + PolynomialNormalForm p2 = PolynomialNormalForm $ Data.Map.filter (/= 0) $ unionWith (+) p1 p2
  -- This is the "FOIL" rule, written as a computer program.
  PolynomialNormalForm p1 * PolynomialNormalForm p2 = sum [PolynomialNormalForm $ singleton (m1 <> m2) (r1 * r2) | (m1,r1) <- assocs p1, (m2,r2) <- assocs p2]
  negate (PolynomialNormalForm p) = PolynomialNormalForm $ fmap negate p
  abs (PolynomialNormalForm p) = PolynomialNormalForm $ fmap abs p
  signum (PolynomialNormalForm p) = PolynomialNormalForm $ fmap signum p
  fromInteger n = case fromInteger n of
    0 -> PolynomialNormalForm mempty
    n' -> PolynomialNormalForm $ singleton mempty n'

-- The set of polynomials is an `r`-algebra, and so their normal-form encoding should be, too.
-- Again, we need to be careful to normalize appropriately.
instance (Num r, Eq r, Ord x) => Algebra r (PolynomialNormalForm r x) where
  scale 0 _ = 0
  scale r (PolynomialNormalForm p) = PolynomialNormalForm . Data.Map.filter (/= 0) $ Data.Map.map (r *) p

-- We need a normal form for polynomials so that we can display them as text in a standard way.
instance (Show r, Show x) => Show (PolynomialNormalForm r x) where
  show (PolynomialNormalForm p) = intercalate " + " $ fmap (show . snd <> show . fst) $ assocs p

-- We display a polynomial by putting it into normal form.
instance (Num r, Eq r, Show r, Ord x, Show x) => Show (Polynomial r x) where
  show = show . eval pnf

-- We compare polynomials for equality by putting them into their respective normal forms and comparing.
instance (Num r, Eq r, Ord x) => Eq (Polynomial r x) where
  p1 == p2 = eval pnf p1 == eval pnf p2

instance (Num r, Ord r, Ord x) => Ord (Polynomial r x) where
  p1 `compare` p2 = eval pnf p1 `compare` eval pnf p2

-- `Var` exists so that we can use strings as variables, but `Var` has a special
-- `Show` instance that will prevent our text from getting cluttered up by quotation marks.
newtype Var = V String
  deriving (Eq, Ord)

-- The `Show` instance for `String` gives the string escaped and quoted.
-- The `Show` instance for `Var` just gives the naked string.
instance Show Var where
  show (V n) = n

-- This lets us construct a value of type `Var` using ordinary `String` syntax.
-- E.g., `"x"` is polymorphic, and it can mean `"x" :: String` or `"x" :: Var`, thanks to this instance.
instance IsString Var where
  fromString = V

x :: Num r => Polynomial r Var
x = var "x"

y :: Num r => Polynomial r Var
y = var "y"

-- 2x^2 + y + 2
p1 :: Num r => Polynomial r Var
p1 = 2 * x ^ 2 + y + 2

-- 2y^3 + 4x^2 + 3y^2 + 3
p2 :: Num r => Polynomial r Var
p2 = 2 * y ^ 3 + 4 * x ^ 2 + 3 * y ^ 2 + 3

p :: Num r => Polynomial r Var
p = p1 + p2

main :: IO ()
main = print (p @Integer)