lol.hs

{-# LANGUAGE GADTs, RankNTypes, TypeFamilies, DataKinds, DeriveFunctor, TypeOperators  #-}

import Data.Type.Equality
import Data.List (intercalate)
import Data.Maybe (catMaybes)

data HasVars = Var | NoVar

data SHasVars a where
  SVar :: SHasVars Var
  SNoVar :: SHasVars NoVar

type family Or (a :: HasVars) (b :: HasVars) :: HasVars where
  Or NoVar NoVar = NoVar
  Or a     b     = Var
  
sor :: SHasVars a -> SHasVars b -> SHasVars (Or a b)
sor SNoVar SNoVar = SNoVar
sor SNoVar SVar = SVar
sor SVar SNoVar = SVar
sor SVar SVar = SVar

data Expr a where
  VarX :: Expr Var
  Lit :: Int -> Expr NoVar
  Add :: Expr a -> Expr b -> Expr (Or a b)
  Times :: Expr a -> Expr b -> Expr (Or a b)
  Exp :: Expr a -> Expr NoVar -> Expr a
  
exprType :: Expr a -> SHasVars a
exprType VarX = SVar
exprType (Lit _) = SNoVar
exprType (Add e1 e2) = sor (exprType e1) (exprType e2)
exprType (Times e1 e2) = sor (exprType e1) (exprType e2)
exprType (Exp e _) = exprType e

{-
q1 :: SHasVars a -> SHasVars b -> (Or a b :~: NoVar) -> (a :~: NoVar)
q1 SNoVar SNoVar Refl = Refl

q2 :: SHasVars a -> SHasVars b -> (Or a b :~: NoVar) -> (b :~: NoVar)
q2 SNoVar SNoVar Refl = Refl
-}

q1' :: (Or a b ~ NoVar) => SHasVars a -> SHasVars b -> (a :~: NoVar)
q1' SNoVar SNoVar = Refl

q2' :: (Or a b ~ NoVar) => SHasVars a -> SHasVars b -> (b :~: NoVar)
q2' SNoVar SNoVar = Refl

{-
evalN :: Expr NoVar -> Int
evalN (Lit n) = n
evalN (Add e e') = (gcastWith (q1 (exprType e) (exprType e') Refl) (evalN e)) + (gcastWith (q2 (exprType e) (exprType e') Refl) (evalN e'))
evalN (Times e e') = (gcastWith (q1 (exprType e) (exprType e') Refl) (evalN e)) * (gcastWith (q2 (exprType e) (exprType e') Refl) (evalN e'))
-}

evalN' :: Expr NoVar -> Int
evalN' (Lit n) = n
evalN' (Add e e') = (gcastWith (q1' (exprType e) (exprType e')) (evalN' e)) + (gcastWith (q2' (exprType e) (exprType e')) (evalN' e'))
evalN' (Times e e') = (gcastWith (q1' (exprType e) (exprType e')) (evalN' e)) * (gcastWith (q2' (exprType e) (exprType e')) (evalN' e'))
evalN' (Exp e e') = (gcastWith (q1' (exprType e) (exprType e')) (evalN' e)) ^ (evalN' e')

newtype Polynomial a = P [a] deriving (Show, Eq, Functor)

pcons :: Num a => a -> Polynomial a -> Polynomial a
pcons x (P xs) = P (x:xs)

paddedZipSum (x:xs) (y:ys) = (x + y) : paddedZipSum xs ys
paddedZipSum xs [] = xs
paddedZipSum [] ys = ys

instance Num a => Num (Polynomial a) where
  fromInteger n = P [fromInteger n]
  
  (P xs) + (P ys) = P $ reverse $ paddedZipSum (reverse xs) (reverse ys)
  
  -- this is probably broken
  (P (x:xs)) * (P (y:ys)) = (x*y) `pcons` ((P [x]) * (P ys) + (P xs) * (P (y:ys)))
  _ * _ = P []

eval :: forall a. Expr a -> Polynomial Int
eval VarX = P [1,0]
eval (Lit n) = P [0,n]
eval (Add e e') = (eval e) + (eval e')
eval (Times e e') = (eval e) * (eval e')
eval (Exp e e') = (eval e) ^ (evalN e')

printTerm :: Int -> Int -> Maybe String
printTerm _ 0 = Nothing
printTerm 0 x = Just $ show x
printTerm 1 x = Just $ concat [show x, "x"]
printTerm n x = Just $ concat [show x, "x^", show n]

printPoly :: Polynomial Int -> String
printPoly (P xs) = intercalate " + " . catMaybes . terms $ xs
  where
    terms = reverse . zipWith printTerm [0..] . reverse