simplify.hs

import Data.Function (on)
import Data.List (groupBy, sortOn)
import Data.Maybe (mapMaybe)

data Expr
  = Var String
  | Const Integer
  | Add Expr Expr
  | Mul Expr Expr
  deriving (Show)

type NormExpr = [(Integer, [String])]

normalize :: Expr -> NormExpr
normalize (Var v) = [(1, [v])]
normalize (Const x) = [(x, [])]
normalize (Add a b) = normalize a <> normalize b
normalize (Mul a b) = [(a_k * b_k, a_vars <> b_vars) | (a_k, a_vars) <- normalize a, (b_k, b_vars) <- normalize b]

simplifyNorm :: NormExpr -> NormExpr
simplifyNorm = fmap (\terms -> (sum (fmap fst terms), snd (head terms))) . groupBy ((==) `on` snd) . sortOn snd

toExpr :: NormExpr -> Expr
toExpr = foldr1 Add . mapMaybe f
  where
    f (0, _) = Nothing
    f (1, []) = Just $ Const 1
    f (1, vars) = Just $ foldl1 Mul (Var <$> vars)
    f (k, vars) = Just $ foldl Mul (Const k) (Var <$> vars)

simplify :: Expr -> Expr
simplify = toExpr . simplifyNorm . normalize