{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts, TypeFamilies, ScopedTypeVariables #-}

-- An attempt to use Map for memoizing functions, including the functions
-- with the infinite domain: see fromMemo and toMemo below.

-- We can memoize not only functions from [Bool] (the infinite domain
-- of all finite boolean lists) but also functions from [[Bool]] 
-- (arbitrary lists of arbitrary lists of booleans). Alas, although
-- the memoized function does terminate (provided that un-memoized does),
-- it may take impractically long.

-- The key is to be able to `fairly' flatten a possibly infinite list of
-- possibly infinite lists. Any element in the argument list
-- shall be found in finite time in the result list. Alas, this finite
-- time may be fairly long (for the case of [[Bool]])

-- The only use of this file is to demonstrate that insertion in the map
-- can be done lazily, and we can really lazily build a Map with the infinite
-- number of keys and look up a key in such a map (see test_fbooll).


module Map where
import Prelude hiding ( lookup )
import qualified Data.IntMap

import qualified Data.List as L

-- Start class
class Key k where
  data Map k :: * -> *
  empty  :: Map k v
  lookup :: k -> Map k v -> Maybe v
  insert :: k -> v -> Map k v -> Map k v
  dom    :: [k]		    -- a potentially infinite domain of all keys
  -- ...many other methods could be added...
-- Stop class

-- Start instances
instance Key Bool where
  data Map Bool elt = MB (Maybe elt) (Maybe elt)
  empty = MB Nothing Nothing
  lookup False (MB mf _) = mf
  lookup True  (MB _ mt) = mt
  insert False v ~(MB _ mt) = MB (Just v) mt
  insert True  v ~(MB mf _) = MB mf (Just v)
  dom = [False, True]

instance (Key a, Key b) => Key (Either a b) where
  data Map (Either a b) elt = MS (Map a elt) (Map b elt)
  empty = MS empty empty
  lookup (Left  k) (MS m _) = lookup k m
  lookup (Right k) (MS _ m) = lookup k m
  insert (Left  k) v (MS ml mr) = MS (insert k v ml) mr
  insert (Right k) v (MS ml mr) = MS ml (insert k v mr)
  dom = interleave (map Left dom) (map Right dom)
  
instance (Key a, Key b) => Key (a,b) where
  data Map (a,b) elt = MP (Map a (Map b elt))
  empty = MP empty
  lookup (a,b) (MP m) = case lookup a m of
			Nothing -> Nothing
			Just m' -> lookup b m'
  insert (a,b) v (MP m) = MP (insert a (insert b v mb) m)
   where mb = maybe empty id $ lookup a m
  dom = fair_interleave $ map (\h -> map ((,) h) dom) dom
-- Stop instances

interleave :: [a] -> [a] -> [a]
interleave [] x = x
interleave x [] = x
interleave (h1:t1) (h2:t2) = h1 : h2 : interleave t1 t2

-- Flatten a possibly infinite list of possibly infinite lists
fair_interleave :: [[a]] -> [a]
fair_interleave l = step [] [] l
 where step [] [] [] = []
       step [] q []  = step q [] []
       step [] q (h1:h2:t) = step (reverse (h1:h2:q)) [] t
       step [] q (h:t) = step (reverse (h:q)) [] t
       step ([]:rest) q l = step rest q l
       step ((h:t):rest) q l = h : step rest (t:q) l
testi = take 10 $ fair_interleave (repeat [1..])
-- [1,1,2,1,3,2,1,3,4,2]

-- Start int
instance Key Int where
  newtype Map Int elt = MI (Data.IntMap.IntMap elt)
  empty = MI Data.IntMap.empty
  lookup k (MI m) = Data.IntMap.lookup k m
  insert k v (MI m) = MI (Data.IntMap.insert k v m)
  dom = 0 : (interleave [1..] [-1,-2..])
-- Stop int

instance (Key a) => Key [a] where
  data Map [a] elt = ML (Maybe elt) (Map a (Map [a] elt))
  empty = ML Nothing empty
  lookup [] ~(ML m0 _) = m0
  lookup (h:t) ~(ML _ m1) = case lookup h m1 of
			   Nothing -> Nothing
			   Just m' -> lookup t m'
  insert [] v ~(ML _ m1) = ML (Just v) m1
  insert (h:t) v ~(ML m0 m1) = ML m0 (insert h (insert t v mb) m1)
   where mb = maybe empty id $ lookup h m1
  dom = fair_interleave (grow [[]])
     where 
     grow l = l : grow (concatMap (\e -> map (e:) l) doma)
     doma = dom

-- Testing construction of infinite domains
tdom1 = take 5 (dom :: [(Bool,Int)])
-- [(False,0),(True,0),(False,1),(False,-1),(True,1)]
tdom2 = take 5 (dom :: [(Int,Bool)])
-- [(0,False),(1,False),(0,True),(-1,False),(1,True)]
tdom3 = take 9 (dom :: [(Int,Int)])
-- [(0,0),(1,0),(0,1),(-1,0),(0,-1),(1,1),(2,0),(1,-1),(0,2)]

tdom41 = take 9 (dom :: [Bool])
-- [False,True]
tdom42 = take 12 (dom :: [[Bool]])
{-
  [[],[False],[False,False],[True],[False,False,False],[False,True],
   [False,False,False,False],[True,False],[False,False,True],
   [False,False,False,False,False],[False,True,False],[True,True]]
-}
tdom43 = take 11 (dom :: [[[Bool]]])

tdom51 = take 11 (dom :: [[[Int]]])

-- Using Map to memoize functions
fromMemo :: Key k => Map k v -> (k -> v)
fromMemo m k = maybe (error "should have been memoized") id $ lookup k m

toMemo :: (Key k) => (k->v) -> Map k v
toMemo f = foldr (\e m -> insert e (f e) m) empty $ dom

-- Testing memoization

-- A sample function to memoize. Its domain is infinite: the set
-- of all finite boolean lists
fbooll :: [Bool] -> Integer
fbooll [] = 0
fbooll (h:t) = (fromIntegral . fromEnum $ h) + 2 * fbooll t

-- Test constructing the Map that maps all boolean lists, the infinite
-- number of them
test_fbooll = (fr,gr)
  where fr = fbooll [True,False,True]
	gr = gbooll [True,False,True]
	gbooll = fromMemo (toMemo fbooll)
-- (5,5)

-- Alas, although the following works in principle, it works in practice
-- only for very small arguments. It just takes really long time
-- for most arguments
test_fbooll2 = (fr,gr)
  where f1 = sum :: [Integer] -> Integer
	f = f1 . (map fbooll)
	-- The commented out sample is too large to handle
	-- sample = replicate 5 [True,False,True]
	sample = replicate 5 [False]
	fr = f sample
	gr = g sample
	g = fromMemo (toMemo f)
-- (0,0)

{- Data.IntMap is not lazy enough ...
test_fint = (fr,gr)
  where f = sum :: [Int] -> Int
	sample = [1..10]
	fr = f sample
	gr = g sample
	g = fromMemo (toMemo f)

-}