IRTreeSpec.hs

{-# LANGUAGE BlockArguments #-}

module Control.Monad.IRTreeSpec where

import Control.Monad.IRTree
import qualified Control.Monad.IRTree as IRTree
import qualified Control.Monad.RTree as RTree
import Control.Monad.Writer.Strict
import Data.Bool
import Data.Expr as Expr
import Data.List.NonEmpty (nonEmpty)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as Map
import Test.Hspec

spec :: Spec
spec = describe "examples" do
  handle "small-opr-expr" $
    Opr (Cnt 1) (Cnt 2)

  handle "small-let-expr" $
    Let "x" (Cnt 1) (Opr (Cnt 2) (Var "x"))

  handle "double-let-expr" $
    Let "x" (Cnt 1) (Let "y" (Cnt 2) (Opr (Var "x") (Var "y")))

  handle "double-overloading-let-expr" $
    Let "x" (Cnt 1) (Let "x" (Cnt 2) (Opr (Var "x") (Var "x")))
 where
  handle str e = describe (str <> " (" <> prettyExprS 0 e ")") do
    let me = runReaderT (Expr.rExpr e) Map.empty

    it "should extract" do
      IRTree.extract me `shouldBe` e

    let re = runReaderT (Expr.rExpr e) Map.empty

    let
      predicate :: Expr -> IO Bool
      predicate = pure . contains isOpr

    rex <- runIO $ RTree.reduce predicate re

    -- onGlitterWith
    --   ("test/expected/" <> str <> "-red")
    --   ( \fp () -> do
    --       (mex, result) <- runWriterT (IRTree.reduce (debugPredicate showString (prettyExprS 0) predicate) me)
    --       writeFile fp (appEndo result "")
    --       pure mex
    --   )
    --   do
    --     it "should produce the same results as the RTree" \mex -> do
    --       rex `shouldBe` mex

    it "should find an opr exponentially" do
      (mex, _) <- runWriterT (IRTree.reduceExp (debugPredicate showString (prettyExprS 0) predicate) me)
      rex `shouldBe` mex

    it "should find an opr fibonacci" do
      (mex, _) <- runWriterT (IRTree.reduceFib (debugPredicate showString (prettyExprS 0) predicate) me)
      rex `shouldBe` mex

    it "should reduce like iinputs" do
      forM_ (RTree.iinputs re) \(ii, e') -> do
        p <- toPredicate ii
        IRTree.reduce (const p) me `shouldReturn` Just e'

debugPredicate
  :: (Monad m)
  => (l -> ShowS)
  -> (i -> ShowS)
  -> (i -> m Bool)
  -> [(Bool, l)]
  -> i
  -> WriterT (Endo String) m Bool
debugPredicate ppl ppi predicate lst i = do
  x <- lift (predicate i)
  tell . Endo $
    showString (bool "0" "1" x)
      . showString ": "
      . ppi i
      . showString "\n"
      . case nonEmpty lst of
        Nothing -> showString "initial\n"
        Just lst' -> ppl (snd $ NE.last lst') . showString "\n"

  pure x