{-# LANGUAGE TypeOperators #-}

module ExprF (
    ExprF(..), Expr(..),
    exprId, exprEval, exprEval'
  ) where

import FixFInstance ()

import Control.Functor.Algebra
import Control.Functor.Fix
import Control.Functor.Composition
import Control.Monad.Either ()

import Prelude hiding (mapM, concatMap)
import Control.Applicative
import Data.Foldable
import Data.Traversable


-- ExprF

-- | The type of arithmetic expressions.
data ExprF r
  = Add r r
  | Sub r r
  | Mul r r
  | Div r r
  | Num Int
  deriving (Eq, Show)

instance Functor ExprF where
  fmap = fmapDefault

instance Foldable ExprF where
  foldMap = foldMapDefault

instance Traversable ExprF where
  traverse f expr = case expr of
    Add x y -> Add <$> f x <*> f y
    Sub x y -> Sub <$> f x <*> f y
    Mul x y -> Mul <$> f x <*> f y
    Div x y -> Div <$> f x <*> f y
    Num n   -> pure (Num n)


-- Fix ExprF as Num instance

-- | Wrapper for 'FixF' 'ExprF', allowing easy construction of expressions:
-- 
-- @> runExpr (4 + 5)
--InF (Add (InF (Num 4)) (InF (Num 5)))@

newtype Expr = Expr { runExpr :: FixF ExprF }
  deriving (Eq, Show)

instance Num Expr where
  fromInteger = Expr . InF . Num . fromIntegral
  Expr x + Expr y = Expr $ InF $ Add x y
  Expr x - Expr y = Expr $ InF $ Sub x y
  Expr x * Expr y = Expr $ InF $ Mul x y
  negate = (0 -)
  abs = error "abs"
  signum = error "signum"

instance Fractional Expr where
  Expr x / Expr y = Expr $ InF $ Div x y

type ErrorAlgebra f e a = f a -> Either e a

-- | The identity algebra.
exprId :: ErrorAlgebra ExprF e (FixF ExprF)
exprId = Right . InF

-- | An algebra that forbids division by zero.
exprEval :: ErrorAlgebra ExprF String Int
exprEval expr = case expr of
  Num n   -> Right n
  Add x y -> Right (x + y)
  Sub x y -> Right (x - y)
  Mul x y -> Right (x * y)
  Div x y
    | y == 0    -> Left "division by zero"
    | otherwise -> Right (x `div` y)

exprEval' :: Algebra ((,) x :.: ExprF) (Either (x, String) Int)
exprEval' expr = case expr of
    CompF (_, Num n)   -> pure n
    CompF (_, Add x y) -> (+) <$> x <*> y
    CompF (_, Sub x y) -> (-) <$> x <*> y
    CompF (_, Mul x y) -> (*) <$> x <*> y
    CompF (z, Div x y) -> do
      x' <- x
      y' <- y
      if y' == 0
        then Left (z, "division by zero")
        else pure (x' `div` y')