module Shared (Shared
              ,extendShared -- Incrementally build dictionary
              ,nullShared   -- Empty dictionary
              ,nameShared   -- Dictionary lookup
              ,foldShared   -- Traverse unique nodes
              ) where

-- Recover shared nodes from a tree.  This is useful for common
-- subexpression elimination relative to some externally defined
-- equivalence relation (Eq t).

-- This code uses memoization of (==) based on makeStableName to
-- reduce the algorithmic complexity of repeated tree comparison.

-- An additional problem this solves is to recover sharing for
-- abstract interpreted code.  Since these will already (likely?) be
-- pointer-equal, this recovery is cheap.


import Control.Monad
import System.IO.Unsafe
import System.Mem.StableName


-- import Memo

-- A Shared environment is constructed using the extendShared function
-- which takes a term and an environment, and returs a shared term and
-- an updated environment.

-- The implementation uses memoization of the membership function for
-- t, based on a table of terms and a table of cached term
-- equivalences.

data Shared term = Shared [term] [(term,term)]

-- To extend an environment of shared terms with a Term is to either
-- establish its membership and return a representative, or to add it
-- to the node or equivalence tables.

extendShared :: Eq t => t -> Shared t -> (t, Shared t)
extendShared t e@(Shared ts eqs) = 
    case (cache0 `mplus` cache1) of
      Just t' -> (t',e)                                       -- memo hit
      Nothing -> case firstWith (t ==) ts of
                   Just t' -> (t', (Shared ts ((t,t'):eqs)))  -- extend memo table
                   Nothing -> (t,  (Shared (t:ts) eqs))       -- extend node table
    where
      cache0 = firstWith  (t `unsafeEq`) ts      -- try exact== node table 
      cache1 = lookupWith (t `unsafeEq`) eqs     -- try equivalence== memo table


-- Once constructed, a Shared environment can be used in terms of a
-- fold over the shared nodes it found..

foldShared :: (t -> a -> a) -> a -> Shared t -> a
foldShared fn i (Shared ts _) = foldr fn i ts

-- .. or to construct a lookup method that translates a shared node to
-- a node sequence number.

-- FIXME: the memoization probably doesn't work here leading to
-- greater computational complexity.  Abstract it properly and make it
-- work.

nameShared :: (Eq term) => 
              [(term,name)] -> [name] -> Shared term -> term -> name
           
nameShared outDict tmpNames (Shared ts _) = name where
    name t = found $ lookupWith (t .==) tagged
    tagged = f (reverse ts) tmpNames where
        f [] _ = []
        f (t:ts) ns@(n:ns') =
            case lookupWith (t .==) outDict of
              Just o  -> (t, o) : f ts ns
              Nothing -> (t, n) : f ts ns'
    t1 .== t2 = t1 `unsafeEq` t2 || t1 == t2

-- Unsafe ops (not referentially transparent).  They always need to be
-- combined with a fallback to functions based on (==) to maintain
-- referential transparency.

unsafeEq a b = unsafePerformIO $ do
                 a' <- makeStableName a
                 b' <- makeStableName b
                 return (a' == b')






-- Tools

instance Show t => Show (Shared t) where
    show (Shared ts eqs) = show (ts,eqs)

found (Just x) = x
found Nothing = error "non-shared term"

nullShared = Shared [] []

firstWith p = f where
    f []           = Nothing
    f (x:xs) | p x = Just x
    f (_:xs)       = f xs

lookupWith p = f where
    f []               = Nothing
    f ((a,b):ps) | p a = Just b
    f (_:ps)           = f ps