-- Preventing memoization in (AI) search problems and
-- non-deterministic programming. 
-- Or, how to build search trees and _avoid_ memoization
-- implicit in lazy evaluation.
-- It is trickier than appears.

module STrees where

import Control.Monad
import System.Environment (getArgs)

-- Sample non-deterministic problem: computing
-- an infinite stream of Pythagorean triples

-- An infinite stream of integers
from :: MonadPlus m => Int -> m Int
from i = return i `mplus` from (i+1)

pyth :: MonadPlus m => m (Int,Int,Int)
pyth = do
  x <- from 1
  y <- from 1
  z <- from 1
  if x*x + y*y == z*z then return (x,y,z) else mzero


-- We represent the result of a non-deterministic computation
-- as an (infinite) search tree.

-- First version of the search tree
-- One may say it is a free MonadPlus
data Tree1 a = Fail1 | Val1 a | Node1 (Tree1 a) (Tree1 a)

instance Monad Tree1 where
    return = Val1

    Fail1  >>= f = Fail1
    Val1 x >>= f = f x
    Node1 e1 e2 >>= f = Node1 (e1 >>= f) (e2 >>= f)

instance MonadPlus Tree1 where
    mzero = Fail1
    mplus = Node1


-- Extracting the results from the search tree as 
-- a lazy list (stream)
-- The depth-first search is hopeless here

-- Breadth-first search
bfs1 :: Tree1 a -> [a]
bfs1 t = loop [t]
 where
 loop [] = []
 loop (Fail1:ts)  = loop ts
 loop (Val1 x:ts) = x : loop ts
 loop (Node1 e1 e2:ts) = loop (ts ++ [e1,e2])

testb1 = take 5 . bfs1 $ pyth
-- [(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13)]

-- Iterative deepening
-- See Searches.hs for more explanation

iter_deepening1 :: Tree1 a -> [a]
iter_deepening1 t = loop 0
 where
 loop d = check d (depth_search1 d t)
 check d Nothing  = []
 check d (Just l) = l ++ loop (d+1)

depth_search1 :: Int -> Tree1 a -> Maybe [a]
depth_search1 d Fail1    = Nothing
depth_search1 0 (Val1 x) = Just [x]
depth_search1 d (Val1 _) = Nothing
depth_search1 0 _        = Just []
depth_search1 d (Node1 e1 e2) = 
  joinM (depth_search1 (d-1) e1) (depth_search1 (d-1) e2)

joinM :: Maybe [a] -> Maybe [a] -> Maybe [a]
joinM x Nothing = x
joinM Nothing x = x
joinM (Just l1) (Just l2) = Just (l1 ++ l2)

testi1 = take 5 . iter_deepening1 $ pyth
-- [(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13)]


-- Compiling the benchmark:
-- ghc -O2 -rtsopts -main-is STrees.main1 STrees.hs
-- To run this code
-- GHCRTS="-tstderr" ./STrees bfs 30
-- GHCRTS="-tstderr" ./STrees iter 30

main0 bfs iter_deepening = getArgs >>= check
 where
 check [key,ns] | [(n,"")] <- reads ns = 
   print $ take n . select key $ pyth
 select "bfs"  = bfs
 select "iter" = iter_deepening

main1 = main0 bfs1 iter_deepening1
{-
./STrees bfs 30 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13),(12,5,13),(9,12,15),(12,9,15),(8,15,17),(15,8,17),(12,16,20),(16,12,20),(7,24,25),(24,7,25),(10,24,26),(15,20,25),(20,15,25),(24,10,26),(20,21,29),(21,20,29),(18,24,30),(24,18,30),(16,30,34),(30,16,34),(12,35,37),(21,28,35),(28,21,35),(35,12,37),(9,40,41),(15,36,39)]
<<ghc: 16091359676 bytes, 30692 GCs, 306378/476332 avg/max bytes residency (12 samples), 3M in use, 0.00 INIT (0.00 elapsed), 12.99 MUT (13.35 elapsed), 5.06 GC (5.22 elapsed) :ghc>>

./STrees iter 30 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13),(12,5,13),(9,12,15),(12,9,15),(8,15,17),(15,8,17),(12,16,20),(16,12,20),(7,24,25),(24,7,25),(10,24,26),(15,20,25),(20,15,25),(24,10,26),(20,21,29),(21,20,29),(18,24,30),(24,18,30),(16,30,34),(30,16,34),(12,35,37),(21,28,35),(28,21,35),(35,12,37),(9,40,41),(15,36,39)]
<<ghc: 53635420 bytes, 103 GCs, 847726/1543008 avg/max bytes residency (6 samples), 5M in use, 0.00 INIT (0.00 elapsed), 0.13 MUT (0.13 elapsed), 0.06 GC (0.07 elapsed) :ghc>>

./STrees iter 100 +RTS -tstderr 
[(3,4,5),... ,(28,96,100),(96,28,100)]
<<ghc: 1825314692 bytes, 3482 GCs, 6100615/21332568 avg/max bytes residency (19 samples), 56M in use, 0.00 INIT (0.00 elapsed), 4.81 MUT (4.90 elapsed), 1.25 GC (1.32 elapsed) :ghc>>

-}


-- Second version of the search tree
-- Attempting to prevent memoization by using explicit thunks
data Tree2 a = Fail2 | Val2 a | Node2 (() -> Tree2 a) (() -> Tree2 a)

instance Monad Tree2 where
    return = Val2

    Fail2  >>= f = Fail2
    Val2 x >>= f = f x
    Node2 e1 e2 >>= f = 
	Node2 (\() -> e1 () >>= f) (\() -> e2 () >>= f)

instance MonadPlus Tree2 where
    mzero = Fail2
    mplus e1 e2 = Node2 (\() -> e1) (\() -> e2)

bfs2 :: Tree2 a -> [a]
bfs2 t = loop [\() -> t]
 where
 loop [] = []
 loop (h:ts) = case h () of
	        Fail2  -> loop ts
                Val2 x -> x : loop ts
		Node2 e1 e2 -> loop (ts ++ [e1,e2])

testb2 = take 5 . bfs2 $ pyth
-- [(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13)]

depth_search2 :: Int -> Tree2 a -> Maybe [a]
depth_search2 _ Fail2    = Nothing
depth_search2 0 (Val2 x) = Just [x]
depth_search2 d (Val2 _) = Nothing
depth_search2 0 _        = Just []
depth_search2 d (Node2 e1 e2) = 
  joinM (depth_search2 (d-1) (e1 ())) (depth_search2 (d-1) (e2 ()))

iter_deepening2 t = loop 0
 where
 loop d = check d (depth_search2 d t)
 check d Nothing = []
 check d (Just l) = l ++ loop (d+1)

testi2 = take 5 . iter_deepening2 $ pyth
-- [(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13)]

main2 = main0 bfs2 iter_deepening2

{-
./STrees bfs 30 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13),(12,5,13),(9,12,15),(12,9,15),(8,15,17),(15,8,17),(12,16,20),(16,12,20),(7,24,25),(24,7,25),(10,24,26),(15,20,25),(20,15,25),(24,10,26),(20,21,29),(21,20,29),(18,24,30),(24,18,30),(16,30,34),(30,16,34),(12,35,37),(21,28,35),(28,21,35),(35,12,37),(9,40,41),(15,36,39)]
<<ghc: 16091380000 bytes, 30693 GCs, 347601/521564 avg/max bytes residency (15 samples), 3M in use, 0.00 INIT (0.00 elapsed), 12.66 MUT (13.05 elapsed), 5.04 GC (5.22 elapsed) :ghc>>

./STrees iter 30 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13),(12,5,13),(9,12,15),(12,9,15),(8,15,17),(15,8,17),(12,16,20),(16,12,20),(7,24,25),(24,7,25),(10,24,26),(15,20,25),(20,15,25),(24,10,26),(20,21,29),(21,20,29),(18,24,30),(24,18,30),(16,30,34),(30,16,34),(12,35,37),(21,28,35),(28,21,35),(35,12,37),(9,40,41),(15,36,39)]
<<ghc: 100502800 bytes, 192 GCs, 1425926/3034900 avg/max bytes residency (6 samples), 8M in use, 0.00 INIT (0.00 elapsed), 0.28 MUT (0.29 elapsed), 0.10 GC (0.10 elapsed) :ghc>>

./STrees iter 100 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),...,(96,28,100)]
<<ghc: 3552618932 bytes, 6777 GCs, 10272972/40186404 avg/max bytes residency (14 samples), 96M in use, 0.00 INIT (0.00 elapsed), 10.59 MUT (10.77 elapsed), 1.66 GC (1.74 elapsed) :ghc>>

-}


-- The third version of the search tree
-- Using thunks and additional trick to prevent the clever GHC memoize
-- even under lambda
data Tree3 a = Fail3 | Val3 a | Node3 (() -> Tree3 a) (() -> Tree3 a)

instance Monad Tree3 where
    return = Val3

    Fail3  >>= f = Fail3
    Val3 x >>= f = f x
    Node3 e1 e2 >>= f = 
	Node3 (app1 e1 f) (app1 e2 f)

instance MonadPlus Tree3 where
    mzero = Fail3
    mplus e1 e2 = Node3 (app e1) (app e2)

-- The functions app and app1 are the trick
{-# NOINLINE app #-}
app e () = e

{-# NOINLINE app1 #-}
app1 e f () = e () >>= f

bfs3 :: Tree3 a -> [a]
bfs3 t = loop [\() -> t]
 where
 loop [] = []
 loop (h:ts) = case h () of
	        Fail3  -> loop ts
                Val3 x -> x : loop ts
		Node3 e1 e2 -> loop (ts ++ [e1,e2])

testb3 = take 5 . bfs3 $ pyth
-- [(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13)]

depth_search3 :: Int -> Tree3 a -> Maybe [a]
depth_search3 _ Fail3    = Nothing
depth_search3 0 (Val3 x) = Just [x]
depth_search3 d (Val3 _) = Nothing
depth_search3 0 _        = Just []
depth_search3 d (Node3 e1 e2) = 
  joinM (depth_search3 (d-1) (e1 ())) (depth_search3 (d-1) (e2 ()))

iter_deepening3 t = loop 0
 where
 loop d = check d (depth_search3 d t)
 check d Nothing = []
 check d (Just l) = l ++ loop (d+1)

testi3 = take 5 . iter_deepening3 $ pyth
-- [(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13)]

main3 = main0 bfs3 iter_deepening3

{-
./STrees bfs 30 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13),(12,5,13),(9,12,15),(12,9,15),(8,15,17),(15,8,17),(12,16,20),(16,12,20),(7,24,25),(24,7,25),(10,24,26),(15,20,25),(20,15,25),(24,10,26),(20,21,29),(21,20,29),(18,24,30),(24,18,30),(16,30,34),(30,16,34),(12,35,37),(21,28,35),(28,21,35),(35,12,37),(9,40,41),(15,36,39)]
<<ghc: 16094522816 bytes, 30699 GCs, 275411/422652 avg/max bytes residency (8 samples), 3M in use, 0.00 INIT (0.00 elapsed), 13.19 MUT (13.81 elapsed), 4.74 GC (5.03 elapsed) :ghc>>

./STrees iter 30 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13),(12,5,13),(9,12,15),(12,9,15),(8,15,17),(15,8,17),(12,16,20),(16,12,20),(7,24,25),(24,7,25),(10,24,26),(15,20,25),(20,15,25),(24,10,26),(20,21,29),(21,20,29),(18,24,30),(24,18,30),(16,30,34),(30,16,34),(12,35,37),(21,28,35),(28,21,35),(35,12,37),(9,40,41),(15,36,39)]
<<ghc: 270730232 bytes, 517 GCs, 54694/80032 avg/max bytes residency (2 samples), 2M in use, 0.00 INIT (0.00 elapsed), 0.38 MUT (0.38 elapsed), 0.03 GC (0.03 elapsed) :ghc>>

./STrees iter 100 +RTS -tstderr 
[(3,4,5),(4,3,5),(6,8,10),(8,6,10),(5,12,13),(12,5,13),(9,12,15),(12,9,15),(8,15,17),(15,8,17),(12,16,20),(16,12,20),(7,24,25),(24,7,25),(10,24,26),(15,20,25),(20,15,25),(24,10,26),(20,21,29),(21,20,29),(18,24,30),(24,18,30),(16,30,34),(30,16,34),(12,35,37),(21,28,35),(28,21,35),(35,12,37),(9,40,41),(15,36,39),(36,15,39),(40,9,41),(24,32,40),(32,24,40),(27,36,45),(36,27,45),(14,48,50),(48,14,50),(20,48,52),(24,45,51),(30,40,50),(40,30,50),(45,24,51),(48,20,52),(28,45,53),(45,28,53),(11,60,61),(33,44,55),(44,33,55),(60,11,61),(40,42,58),(42,40,58),(16,63,65),(36,48,60),(48,36,60),(63,16,65),(25,60,65),(60,25,65),(33,56,65),(56,33,65),(39,52,65),(52,39,65),(32,60,68),(60,32,68),(21,72,75),(24,70,74),(42,56,70),(56,42,70),(70,24,74),(72,21,75),(48,55,73),(55,48,73),(18,80,82),(30,72,78),(45,60,75),(60,45,75),(72,30,78),(80,18,82),(13,84,85),(84,13,85),(48,64,80),(64,48,80),(36,77,85),(77,36,85),(40,75,85),(75,40,85),(51,68,85),(68,51,85),(39,80,89),(80,39,89),(35,84,91),(60,63,87),(63,60,87),(84,35,91),(54,72,90),(72,54,90),(20,99,101),(99,20,101),(28,96,100),(96,28,100)]
<<ghc: 9952569856 bytes, 18984 GCs, 177780/423220 avg/max bytes residency (41 samples), 2M in use, 0.00 INIT (0.00 elapsed), 13.36 MUT (13.93 elapsed), 0.89 GC (0.91 elapsed) :ghc>>

-}