{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
module Main where

import Data.Expr.SharingRecovery

import Arith


-- TODO: test cyclic expressions


a_bin :: (KnownType a, KnownType b, KnownType c)
      => PrimOp (a, b) c
      -> PHOASExpr Typ v ArithF a
      -> PHOASExpr Typ v ArithF b
      -> PHOASExpr Typ v ArithF c
a_bin op a b = PHOASOp τ (A_Prim op (PHOASOp τ (A_Pair a b)))

lam :: (KnownType a, KnownType b)
    => (PHOASExpr Typ v f a -> PHOASExpr Typ v f b) -> PHOASExpr Typ v f (a -> b)
lam f = PHOASLam τ τ $ \arg -> f (PHOASVar τ arg)

(+!) :: PHOASExpr Typ v ArithF Int -> PHOASExpr Typ v ArithF Int -> PHOASExpr Typ v ArithF Int
(+!) = a_bin POAddI

(*!) :: PHOASExpr Typ v ArithF Int -> PHOASExpr Typ v ArithF Int -> PHOASExpr Typ v ArithF Int
(*!) = a_bin POMulI

-- λx. x + x
ea_1 :: PHOASExpr Typ v ArithF (Int -> Int)
ea_1 = lam $ \arg -> arg +! arg

-- λx. let y = x + x in y * y
ea_2 :: PHOASExpr Typ v ArithF (Int -> Int)
ea_2 = lam $ \arg -> let y = arg +! arg
                     in y *! y

ea_3 :: PHOASExpr Typ v ArithF (Int -> Int)
ea_3 = lam $ \arg ->
  let y = arg +! arg
      x = y *! arg
  -- in (y +! x) +! (x +! y)
  in (x +! y) +! (y +! x)

main :: IO ()
main = putStrLn $ prettyBExpr prettyArithF (sharingRecovery ea_3)