# LANGUAGE GADTs , RankNTypes , ScopedTypeVariables #

module Sharing where

import Control.Applicative
import Control.Monad.Fix
import Data.HashTable                           as Hash
import qualified Data.IntMap                    as IntMap
import Debug.Trace
import System.IO.Unsafe                         (unsafePerformIO)
import Data.List
import Data.Maybe
import Data.Typeable
import System.Mem.StableName

import HOAS                                     (Level)
import qualified HOAS as HOAS


-- Occurences maps
-- ---------------

-- Opaque stable name for term nodes. It is being used to key the occurence map; hence, we need to
-- hide the type parameter.
--
data StableTermName where
  StableTermName :: Typeable t => StableName (HOAS.Term t) -> StableTermName

instance Show StableTermName where
  show (StableTermName sn) = show $ hashStableName sn

instance Eq StableTermName where
  StableTermName sn1 == StableTermName sn2
    | Just sn1' <- gcast sn1 = sn1' == sn2
    | otherwise              = False

makeStableTerm :: HOAS.Term t -> IO (StableName (HOAS.Term t))
makeStableTerm e = e `seq` makeStableName e

-- Mutable occurence map

-- Mutable hashtable version of the occurrence map keyed on the stable names of terms. It associates
  each term node with an occurence count and the height of the AST .
--
type OccMapHash = Hash.HashTable StableTermName (Int, Int)

  Create a new hash table keyed on AST nodes .
--
newOccMapHashTable :: IO OccMapHash
newOccMapHashTable = Hash.new (==) hashStableTerm
  where
    hashStableTerm (StableTermName sn) = fromIntegral (hashStableName sn)

  Enter one term node occurrence into an occurrence map .   Returns ' Just h ' if this is a repeated
-- occurence and the height of the repeatedly occuring term is 'h'.
--
  If this is the first occurence , the ' height ' * argument * must provide the height of the term ;
-- otherwise, the height will be *extracted* from the occurence map.  In the latter case, this
-- function yields the term height.
--
enterOcc :: OccMapHash -> StableTermName -> Int -> IO (Maybe Int)
enterOcc occMap sa height
  = do
      entry <- Hash.lookup occMap sa
      case entry of
        Nothing           -> Hash.insert occMap sa (1    , height)  >> return Nothing
        Just (n, heightS) -> Hash.update occMap sa (n + 1, heightS) >> return (Just heightS)    

-- Immutable occurence map

-- Immutable version of the occurence map (storing the occurence count only, not the height).  We
  use the ' StableName ' hash to index an ' IntMap ' and disambiguate ' StableName 's with identical
  hashes explicitly , storing them in a list in the ' IntMap ' .
--
type OccMap = IntMap.IntMap [(StableTermName, Int)]

-- Turn a mutable into an immutable occurence map.
--
freezeOccMap :: OccMapHash -> IO OccMap
freezeOccMap oc
  = do
      kvs <- map dropHeight <$> Hash.toList oc
      return . IntMap.fromList . map (\kvs -> (key (head kvs), kvs)). groupBy sameKey $ kvs
  where
    key (StableTermName sn, _) = hashStableName sn
    sameKey kv1 kv2            = key kv1 == key kv2
    dropHeight (k, (cnt, _))   = (k, cnt)

-- Look up the occurence map keyed by array computations using a stable name.  If a the key does
-- not exist in the map, return an occurence count of '1'.
--
lookupWithTermName :: OccMap -> StableTermName -> Int
lookupWithTermName oc sa@(StableTermName sn) 
  = fromMaybe 1 $ IntMap.lookup (hashStableName sn) oc >>= Prelude.lookup sa

-- Look up the occurence map keyed by array computations using a sharing array computation.  If an
-- the key does not exist in the map, return an occurence count of '1'.
--
lookupWithSharingTerm :: OccMap -> StableSharingTerm -> Int
lookupWithSharingTerm oc (StableSharingTerm (StableTermHeight sn _) _) 
  = lookupWithTermName oc (StableTermName sn)


-- Term structure for sharing recovery
-- -----------------------------------

  Terms for sharing recovery consist of two mutually recursive datatypes . The one hear hold the
-- actual lambda term forms. 'SharingTerm' keeps track of stable names (for identification) as well
-- as where variables and let bindings need to be introduced.
--
data Term binder t where
  for conversion to de Bruijn
  Tag :: Typeable t                               
      => Level                                                -> Term binder t    

  Con :: (Typeable t, Show t)                     
      => t                                                    -> Term binder t
  Lam :: (Typeable s, Typeable t, Show s, Show t) 
      => binder -> SharingTerm binder t                       -> Term binder (s -> t)
  App :: (Typeable s, Typeable t, Show s, Show t) 
      => SharingTerm binder (s -> t) -> SharingTerm binder  s -> Term binder t

showTermOp :: Term binder t -> String
showTermOp (Tag lvl) = "Tag " ++ show lvl
showTermOp (Con v)   = "Con " ++ show v
showTermOp (Lam {})  = "Lam"
showTermOp (App {})  = "App"


-- Stable term nodes
-- -----------------

-- Stable name for a term node including the height of the term.
--
data StableTermHeight t = StableTermHeight (StableName (HOAS.Term t)) Int

instance Eq (StableTermHeight t) where
  (StableTermHeight sn1 _) == (StableTermHeight sn2 _) = sn1 == sn2

higherSTH :: StableTermHeight t1 -> StableTermHeight t2 -> Bool
StableTermHeight _ h1 `higherSTH` StableTermHeight _ h2 = h1 > h2

hashStableTermHeight :: StableTermHeight t -> Int
hashStableTermHeight (StableTermHeight sn _) = hashStableName sn


  Sharing information in terms
-- ----------------------------

  Interleave sharing annotations into a term .    can be marked as being represented by
  variable ( binding a shared subtree ) using ' VarSharing ' and as being prefixed by a let binding
  ( for a shared subtree ) using ' LetSharing ' .
--
data SharingTerm binder t where
  VarSharing  :: Typeable t
              => StableTermHeight t                        -> SharingTerm binder t
  LetSharing  :: StableSharingTerm -> SharingTerm binder t -> SharingTerm binder t
  TermSharing :: Typeable t
              => StableTermHeight t -> Term binder t       -> SharingTerm binder t

-- Stable name for a term associated with its sharing-annotated version.
--
data StableSharingTerm where
  StableSharingTerm :: Typeable t => StableTermHeight t -> SharingTermFloated t -> StableSharingTerm

instance Show StableSharingTerm where
  show (StableSharingTerm sn _) = show $ hashStableTermHeight sn

instance Eq StableSharingTerm where
  StableSharingTerm sn1 _ == StableSharingTerm sn2 _
    | Just sn1' <- gcast sn1 = sn1' == sn2
    | otherwise              = False

higherSST :: StableSharingTerm -> StableSharingTerm -> Bool
StableSharingTerm sn1 _ `higherSST` StableSharingTerm sn2 _ = sn1 `higherSTH` sn2

-- Test whether the given stable names matches an array computation with sharing.
--
matchStableTerm :: Typeable t => StableTermHeight t -> StableSharingTerm -> Bool
matchStableTerm sn1 (StableSharingTerm sn2 _)
  | Just sn1' <- gcast sn1 = sn1' == sn2
  | otherwise              = False

-- Dummy entry for environments to be used for unused variables.
--
noStableTermName :: StableTermHeight t
noStableTermName = unsafePerformIO $ StableTermHeight <$> makeStableName undefined <*> pure 0


  Sharing recovery
-- ----------------

-- Sharing terms after computing the occurence map and pruning repeated subtrees.
--
  Lambdas are annotated with the level of their binder .
--
type SharingTermPruned  t = SharingTerm Level t

-- Sharing terms after sharing recovery and shared subtrees have been floated to their let-binding
-- positions.
--
  Lambdas are annotated by the ' StableSharingTerm ' identifying their binder .
--
type SharingTermFloated t = SharingTerm StableSharingTerm t

-- |Recover sharing information and annotate the HOAS AST with variable and let binding annotations.
--
-- NB: Strictly speaking, this function is not deterministic, as it uses stable pointers to
--     determine the sharing of subterms.  The stable pointer API does not guarantee its
--     completeness; i.e., it may miss some equalities, which implies that we may fail to discover
--     some sharing.  However, sharing does not affect the denotational meaning of a term; hence,
--     we do not compromise denotational correctness.
--
      There is one caveat : We currently rely on ' Tag ' leaves representing free variables to be
--       shared if any of them is used more than once.  If one is duplicated, the environment for
--       de Bruijn conversion will have a duplicate entry, and hence, be of the wrong size, which
--       is fatal. (The 'lookupStableTag' function will already bail out.)
--
recoverSharing :: Typeable t => HOAS.Term t -> SharingTermFloated t
# NOINLINE recoverSharing #
recoverSharing term
  = let (term', occMap) =
          unsafePerformIO $ do        -- to enable stable pointers; it's safe as explained above
          { (term', occMap) <- makeOccMap 0 term
          ; frozenOccMap    <- freezeOccMap occMap
          ; return (term', frozenOccMap)
          }
    in 
    determineScopes occMap term'

-- Compute the term occurence map, mark all nodes with stable names, and drop repeated occurences
  of shared subterms ( Phase One ) .
--
-- Note [Traversing functions and side effects]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- We need to descent into function bodies to build the 'OccMap' with all occurences in the
-- function bodies.  Due to the side effects in the construction of the occurence map and, more
  importantly , the dependence of the second phase on /global/ occurence information , we may not
-- delay the body traversals by putting them under a lambda.  Hence, we apply each function, to
-- traverse its body and use a /dummy abstraction/ of the result.
--
-- For example, given a function 'f', we traverse 'f (Tag 0)', which yields a transformed body 'e'.
-- As the result of the traversal of the overall function, we use 'const e'.  Hence, it is crucial
  that the ' Tag ' supplied during the initial traversal is already the one required by the HOAS to
-- de Bruijn conversion in 'Convert.convertSharing' — any subsequent application of 'const e' will
-- only yield 'e' with the embedded 'Tag 0' of the original application.  During sharing recovery,
-- we float /all/ free variables ('Tag' terms) out to construct the initial environment for producing
  de Bruijn indices , which replaces them by ' VarSharing ' nodes .   Hence , the tag values only serve
  the purpose of determining the ordering in that initial environment . They are /not/ directly used
-- to compute the de Brujin indices.
--
makeOccMap :: Typeable t => Level -> HOAS.Term t -> IO (SharingTermPruned t, OccMapHash)
makeOccMap lvl rootTerm
  = do
    { occMap <- newOccMapHashTable
    ; (rootTerm', _) <- traverse lvl occMap rootTerm
    ; return (rootTerm', occMap)
    }
  where
    traverse :: forall t. Typeable t
                => Level -> OccMapHash -> HOAS.Term t -> IO (SharingTermPruned t, Int)
    traverse lvl occMap term
      = mfix $ \ ~(_, height) -> do 
  Compute stable name and enter it into the occurence map
        ; sn                        <- makeStableTerm term
        ; heightIfRepeatedOccurence <- enterOcc occMap (StableTermName sn) height
          
        ; traceLine (HOAS.showTermOp term) $
            case heightIfRepeatedOccurence of
              Just height -> "REPEATED occurence (sn = " ++ show (hashStableName sn) ++ 
                             "; height = " ++ show height ++ ")"
              Nothing     -> "first occurence (sn = " ++ show (hashStableName sn) ++ ")"

            -- Reconstruct the term in shared form.
            --
            -- In case of a repeated occurence, the height comes from the occurence map; otherwise,
            -- it is computed by the traversal function passed in 'newTerm'.  See also 'enterOcc'.
            --
  NB : This function can only be used in the case alternatives below ; outside of the
      case we can not discharge the '  t ' constraint .
        ; let reconstruct :: Typeable t
                          => IO (Term        Level t, Int)
                          -> IO (SharingTerm Level t, Int)
              reconstruct newTerm
                = case heightIfRepeatedOccurence of 
                    Just height 
                      -> return (VarSharing (StableTermHeight sn height), height)
                    _ -> do
                         { (term, height) <- newTerm
                         ; return (TermSharing (StableTermHeight sn height) term, height)
                         }

        ; case term of
            HOAS.Tag i   -> reconstruct $ return (Tag i, 0)           -- height is 0!
            HOAS.Con v   -> reconstruct $ return (Con v, 1)
            HOAS.Lam f   -> reconstruct $ do
                            {   -- see Note [Traversing functions and side effects]
                            ; (body, h) <- traverse (lvl + 1) occMap (f (HOAS.Tag lvl))
                            ; return (Lam lvl body, h + 1)
                            }
            HOAS.App f a -> reconstruct $ do
                            { (f', h1) <- traverse lvl occMap f
                            ; (a', h2) <- traverse lvl occMap a
                            ; return (App f' a', h1 `max` h2 + 1)
                            }
        }

  Type used to maintain how often each shared subterm , so far , occured during a bottom - up sweep .
--
--   Invariants: 
    - If one shared term 's ' is itself a subterm of another shared term ' t ' , then 's ' must occur
      * after * ' t ' in the ' NodeCounts ' .
--   - No shared term occurs twice.
    - A term may have a final occurence count of only 1 iff it is either a free variable ( ' Tag ' )
--     or an array computation listed out of an expression.
--
-- We determine the subterm property by using the tree height in 'StableTermHeight'.  Trees get
  smaller towards the end of a ' NodeCounts ' list .   The height of free variables ( ' Tag ' ) is 0 ,
  whereas other leaves have height 1 .   This guarantees that all free variables are at the end
  of the ' NodeCounts ' list .
--
-- To ensure the invariant is preserved over merging node counts from sibling subterms, the
-- function '(+++)' must be used.
--
type NodeCounts = [NodeCount]

data NodeCount = NodeCount StableSharingTerm Int
               deriving Show

-- Empty node counts
--
noNodeCounts :: NodeCounts
noNodeCounts = []

  Singleton node counts for ' Acc '
--
nodeCount :: StableSharingTerm -> Int -> NodeCounts
nodeCount sst n = [NodeCount sst n]

  Combine node counts that belong to the same node .
--
-- * We assume that the node counts invariant —subterms follow their parents— holds for both
--   arguments and guarantee that it still holds for the result.
--
(+++) :: NodeCounts -> NodeCounts -> NodeCounts
us +++ vs = foldr insert us vs
  where
    insert x               []                         = [x]
    insert x@(NodeCount st1 count1) ys@(y@(NodeCount st2 count2) : ys') 
      | st1 == st2          = NodeCount (st1 `pickNoneVar` st2) (count1 + count2) : ys'
      | st1 `higherSST` st2 = x : ys
      | otherwise           = y : insert x ys'

    (StableSharingTerm _ (VarSharing _))  `pickNoneVar`  st2  = st2
    st1                                   `pickNoneVar`  _st2 = st1
    
  Determine whether a ' NodeCount ' is for a ' Tag ' , which represent free variables .
--
isFreeVar :: NodeCount -> Bool
isFreeVar (NodeCount (StableSharingTerm _ (TermSharing _ (Tag  _))) _) = True
isFreeVar _                                                            = False

  Determine the scopes of all variables representing shared subterms ( Phase Two ) in a bottom - up
-- sweep.
--
  Precondition : there are only ' VarSharing ' and ' TermSharing ' nodes in the argument .
--
determineScopes :: Typeable t
                => OccMap -> SharingTermPruned t -> SharingTermFloated t
determineScopes occMap rootTerm 
  = let
      (sharingTerm, counts) = scopes rootTerm
    in
    if null counts
    then sharingTerm
    else error $ "determineScopes: unbound shared subtrees: " ++ show counts
  where
    scopes :: forall t. SharingTermPruned t -> (SharingTermFloated t, NodeCounts)
    scopes (LetSharing _ _)
      = error $ "determineScopes: scopesAcc: unexpected 'LetSharing'"
    scopes (VarSharing sn)
      = (VarSharing sn, StableSharingTerm sn (VarSharing sn) `nodeCount` 1)
    scopes (TermSharing sn pterm)
      = case pterm of
          Tag i     -> reconstruct (Tag i) noNodeCounts
          Con v     -> reconstruct (Con v) noNodeCounts
          Lam lvl f -> let
                         (f', count)         = scopes f
                         (stableTag, count') = lookupStableTag lvl count
                       in
                       reconstruct (Lam stableTag f') count'
          App f a   -> let
                         (f', count1) = scopes f
                         (a', count2) = scopes a
                       in 
                       reconstruct (App f' a') (count1 +++ count2)
      where
  Occurence count of the currently processed node
        occCount = let StableTermHeight sn' _ = sn
                   in
                   lookupWithTermName occMap (StableTermName sn')

        -- Reconstruct the current tree node.
        --
  * If the current node is being shared ( ' occCount > 1 ' ) , replace it by a ' VarSharing '
    node and float the shared subtree out wrapped in a ' NodeCounts ' value .
        -- * If the current node is not shared, reconstruct it in place.
        -- * Special case for free variables ('Tag'): Replace the tree by a sharing variable and
    float the ' Tag ' out in a ' NodeCounts ' value .   This is idependent of the number of
        --   occurences.
        --
  In either case , any completed ' NodeCounts ' are injected as bindings using ' LetSharing '
        -- node.
        -- 
        reconstruct :: Typeable t 
                    => Term StableSharingTerm t -> NodeCounts -> (SharingTermFloated t, NodeCounts)
        reconstruct newTerm@(Tag _) _subCount
              -- free variable => replace by a sharing variable regardless of the number of occ.s
          = let thisCount = StableSharingTerm sn (TermSharing sn newTerm) `nodeCount` 1
            in
            tracePure "FREE" (show thisCount) $
            (VarSharing sn, thisCount)
        reconstruct newTerm subCount
              -- shared subtree => replace by a sharing variable
          | occCount > 1
          = let allCount = (StableSharingTerm sn sharingTerm `nodeCount` 1) +++ newCount
            in
            tracePure ("SHARED" ++ completed) (show allCount) $
            (VarSharing sn, allCount)
              -- neither shared nor free variable => leave it as it is
          | otherwise
          = tracePure ("Normal" ++ completed) (show newCount) $
            (sharingTerm, newCount)
          where
              -- Determine the bindings that needs to be attached to the current node...
            (newCount, bindHere) = filterCompleted subCount
    
  ... and wrap them in ' LetSharing ' constructors
            lets        = foldl (flip (.)) id . map LetSharing $ bindHere
            sharingTerm = lets $ TermSharing sn newTerm
    
              -- trace support
            completed | null bindHere = ""
                      | otherwise     = "(" ++ show (length bindHere) ++ " lets)"
    
        -- Extract *leading* nodes that have a complete node count (i.e., their node count is equal
        -- to the number of occurences of that node in the overall expression).
        -- 
        -- Nodes with a completed node count should be let bound at the currently processed node.
        --
  NB : Only extract leading nodes ( i.e. , the longest run at the * front * of the list that is
        --     complete).  Otherwise, we would let-bind subterms before their parents, which leads
        --     to scope errors.
        --
        filterCompleted :: NodeCounts -> (NodeCounts, [StableSharingTerm])
        filterCompleted counts
          = let (completed, counts') = break notComplete counts
            in (counts', [sa | NodeCount sa _ <- completed])
          where
            -- a node is not yet complete while the node count 'n' is below the overall number
            -- of occurences for that node in the whole program, with the exception that free
            -- variables are never complete
            notComplete nc@(NodeCount st n) | not . isFreeVar $ nc = lookupWithSharingTerm occMap st > n
            notComplete _                                          = True

-- Find the stable tag representing the binder at the given level (which equals the tag value).
--
-- If there is not such tag, the binder has no usage occurence.
--
-- if the tag occurs multiple times, sharing of the binder tag was not preserved and we cannot
  continue ( c.f . , comments at ' determineScopes ' ) .
-- 
lookupStableTag :: Level -> NodeCounts -> (StableSharingTerm, NodeCounts)
lookupStableTag lvl counts 
  = case partition hasLevel counts of
      ([], _)                     -> (noStableSharing, counts)  -- tag not used in expression
      ([NodeCount st _], counts') -> (st, counts')              -- tag has a unique occurence
      (counts', _)    -> 
        error $ "lookupStableTag: duplicate 'Tag's\n  " ++ 
                intercalate ", " [showST st | NodeCount st _ <- counts']
  where
    hasLevel (NodeCount (StableSharingTerm _ (TermSharing _ (Tag lvl'))) _) = lvl == lvl'
    hasLevel (NodeCount st _)                                               = False
        
    noStableSharing :: StableSharingTerm
    noStableSharing = StableSharingTerm noStableTermName (undefined :: SharingTermFloated ())

    showST (StableSharingTerm _ (TermSharing  sn term)) = show (hashStableTermHeight sn) ++ ": " ++ 
                                                          showTermOp term
    showST (StableSharingTerm _ (VarSharing sn))        = "VarSharing " ++ show (hashStableTermHeight sn)
    showST (StableSharingTerm _ (LetSharing st _ ))     = "LetSharing " ++ show st ++ "..."


-- Debugging
-- ---------

traceOn :: Bool
traceOn = False

tracePure :: String -> String -> a -> a
tracePure header msg val | traceOn   = trace (header ++ ": " ++ msg) val
                         | otherwise = val

traceLine :: String -> String -> IO ()
traceLine header msg = tracePure header msg $ return ()

 Occurences maps
 ---------------
 Opaque stable name for term nodes. It is being used to key the occurence map; hence, we need to
 hide the type parameter.

 Mutable occurence map
 Mutable hashtable version of the occurrence map keyed on the stable names of terms. It associates


 occurence and the height of the repeatedly occuring term is 'h'.

 otherwise, the height will be *extracted* from the occurence map.  In the latter case, this
 function yields the term height.

 Immutable occurence map
 Immutable version of the occurence map (storing the occurence count only, not the height).  We

 Turn a mutable into an immutable occurence map.

 Look up the occurence map keyed by array computations using a stable name.  If a the key does
 not exist in the map, return an occurence count of '1'.

 Look up the occurence map keyed by array computations using a sharing array computation.  If an
 the key does not exist in the map, return an occurence count of '1'.

 Term structure for sharing recovery
 -----------------------------------
 actual lambda term forms. 'SharingTerm' keeps track of stable names (for identification) as well
 as where variables and let bindings need to be introduced.

 Stable term nodes
 -----------------
 Stable name for a term node including the height of the term.

 ----------------------------

 Stable name for a term associated with its sharing-annotated version.

 Test whether the given stable names matches an array computation with sharing.

 Dummy entry for environments to be used for unused variables.

 ----------------
 Sharing terms after computing the occurence map and pruning repeated subtrees.


 Sharing terms after sharing recovery and shared subtrees have been floated to their let-binding
 positions.


 |Recover sharing information and annotate the HOAS AST with variable and let binding annotations.

 NB: Strictly speaking, this function is not deterministic, as it uses stable pointers to
     determine the sharing of subterms.  The stable pointer API does not guarantee its
     completeness; i.e., it may miss some equalities, which implies that we may fail to discover
     some sharing.  However, sharing does not affect the denotational meaning of a term; hence,
     we do not compromise denotational correctness.

       shared if any of them is used more than once.  If one is duplicated, the environment for
       de Bruijn conversion will have a duplicate entry, and hence, be of the wrong size, which
       is fatal. (The 'lookupStableTag' function will already bail out.)

 to enable stable pointers; it's safe as explained above
 Compute the term occurence map, mark all nodes with stable names, and drop repeated occurences

 Note [Traversing functions and side effects]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We need to descent into function bodies to build the 'OccMap' with all occurences in the
 function bodies.  Due to the side effects in the construction of the occurence map and, more
 delay the body traversals by putting them under a lambda.  Hence, we apply each function, to
 traverse its body and use a /dummy abstraction/ of the result.

 For example, given a function 'f', we traverse 'f (Tag 0)', which yields a transformed body 'e'.
 As the result of the traversal of the overall function, we use 'const e'.  Hence, it is crucial
 de Bruijn conversion in 'Convert.convertSharing' — any subsequent application of 'const e' will
 only yield 'e' with the embedded 'Tag 0' of the original application.  During sharing recovery,
 we float /all/ free variables ('Tag' terms) out to construct the initial environment for producing
 to compute the de Brujin indices.

 Reconstruct the term in shared form.

 In case of a repeated occurence, the height comes from the occurence map; otherwise,
 it is computed by the traversal function passed in 'newTerm'.  See also 'enterOcc'.

 height is 0!
 see Note [Traversing functions and side effects]

   Invariants: 
   - No shared term occurs twice.
     or an array computation listed out of an expression.

 We determine the subterm property by using the tree height in 'StableTermHeight'.  Trees get

 To ensure the invariant is preserved over merging node counts from sibling subterms, the
 function '(+++)' must be used.

 Empty node counts



 * We assume that the node counts invariant —subterms follow their parents— holds for both
   arguments and guarantee that it still holds for the result.


 sweep.


 Reconstruct the current tree node.

 * If the current node is not shared, reconstruct it in place.
 * Special case for free variables ('Tag'): Replace the tree by a sharing variable and
   occurences.

 node.
 
 free variable => replace by a sharing variable regardless of the number of occ.s
 shared subtree => replace by a sharing variable
 neither shared nor free variable => leave it as it is
 Determine the bindings that needs to be attached to the current node...
 trace support
 Extract *leading* nodes that have a complete node count (i.e., their node count is equal
 to the number of occurences of that node in the overall expression).
 
 Nodes with a completed node count should be let bound at the currently processed node.

     complete).  Otherwise, we would let-bind subterms before their parents, which leads
     to scope errors.

 a node is not yet complete while the node count 'n' is below the overall number
 of occurences for that node in the whole program, with the exception that free
 variables are never complete
 Find the stable tag representing the binder at the given level (which equals the tag value).

 If there is not such tag, the binder has no usage occurence.

 if the tag occurs multiple times, sharing of the binder tag was not preserved and we cannot
 
 tag not used in expression
 tag has a unique occurence
 Debugging
 ---------
# LANGUAGE GADTs , RankNTypes , ScopedTypeVariables #

module Sharing where

import Control.Applicative
import Control.Monad.Fix
import Data.HashTable                           as Hash
import qualified Data.IntMap                    as IntMap
import Debug.Trace
import System.IO.Unsafe                         (unsafePerformIO)
import Data.List
import Data.Maybe
import Data.Typeable
import System.Mem.StableName

import HOAS                                     (Level)
import qualified HOAS as HOAS



data StableTermName where
  StableTermName :: Typeable t => StableName (HOAS.Term t) -> StableTermName

instance Show StableTermName where
  show (StableTermName sn) = show $ hashStableName sn

instance Eq StableTermName where
  StableTermName sn1 == StableTermName sn2
    | Just sn1' <- gcast sn1 = sn1' == sn2
    | otherwise              = False

makeStableTerm :: HOAS.Term t -> IO (StableName (HOAS.Term t))
makeStableTerm e = e `seq` makeStableName e


  each term node with an occurence count and the height of the AST .
type OccMapHash = Hash.HashTable StableTermName (Int, Int)

  Create a new hash table keyed on AST nodes .
newOccMapHashTable :: IO OccMapHash
newOccMapHashTable = Hash.new (==) hashStableTerm
  where
    hashStableTerm (StableTermName sn) = fromIntegral (hashStableName sn)

  Enter one term node occurrence into an occurrence map .   Returns ' Just h ' if this is a repeated
  If this is the first occurence , the ' height ' * argument * must provide the height of the term ;
enterOcc :: OccMapHash -> StableTermName -> Int -> IO (Maybe Int)
enterOcc occMap sa height
  = do
      entry <- Hash.lookup occMap sa
      case entry of
        Nothing           -> Hash.insert occMap sa (1    , height)  >> return Nothing
        Just (n, heightS) -> Hash.update occMap sa (n + 1, heightS) >> return (Just heightS)    


  use the ' StableName ' hash to index an ' IntMap ' and disambiguate ' StableName 's with identical
  hashes explicitly , storing them in a list in the ' IntMap ' .
type OccMap = IntMap.IntMap [(StableTermName, Int)]

freezeOccMap :: OccMapHash -> IO OccMap
freezeOccMap oc
  = do
      kvs <- map dropHeight <$> Hash.toList oc
      return . IntMap.fromList . map (\kvs -> (key (head kvs), kvs)). groupBy sameKey $ kvs
  where
    key (StableTermName sn, _) = hashStableName sn
    sameKey kv1 kv2            = key kv1 == key kv2
    dropHeight (k, (cnt, _))   = (k, cnt)

lookupWithTermName :: OccMap -> StableTermName -> Int
lookupWithTermName oc sa@(StableTermName sn) 
  = fromMaybe 1 $ IntMap.lookup (hashStableName sn) oc >>= Prelude.lookup sa

lookupWithSharingTerm :: OccMap -> StableSharingTerm -> Int
lookupWithSharingTerm oc (StableSharingTerm (StableTermHeight sn _) _) 
  = lookupWithTermName oc (StableTermName sn)



  Terms for sharing recovery consist of two mutually recursive datatypes . The one hear hold the
data Term binder t where
  for conversion to de Bruijn
  Tag :: Typeable t                               
      => Level                                                -> Term binder t    

  Con :: (Typeable t, Show t)                     
      => t                                                    -> Term binder t
  Lam :: (Typeable s, Typeable t, Show s, Show t) 
      => binder -> SharingTerm binder t                       -> Term binder (s -> t)
  App :: (Typeable s, Typeable t, Show s, Show t) 
      => SharingTerm binder (s -> t) -> SharingTerm binder  s -> Term binder t

showTermOp :: Term binder t -> String
showTermOp (Tag lvl) = "Tag " ++ show lvl
showTermOp (Con v)   = "Con " ++ show v
showTermOp (Lam {})  = "Lam"
showTermOp (App {})  = "App"



data StableTermHeight t = StableTermHeight (StableName (HOAS.Term t)) Int

instance Eq (StableTermHeight t) where
  (StableTermHeight sn1 _) == (StableTermHeight sn2 _) = sn1 == sn2

higherSTH :: StableTermHeight t1 -> StableTermHeight t2 -> Bool
StableTermHeight _ h1 `higherSTH` StableTermHeight _ h2 = h1 > h2

hashStableTermHeight :: StableTermHeight t -> Int
hashStableTermHeight (StableTermHeight sn _) = hashStableName sn


  Sharing information in terms

  Interleave sharing annotations into a term .    can be marked as being represented by
  variable ( binding a shared subtree ) using ' VarSharing ' and as being prefixed by a let binding
  ( for a shared subtree ) using ' LetSharing ' .
data SharingTerm binder t where
  VarSharing  :: Typeable t
              => StableTermHeight t                        -> SharingTerm binder t
  LetSharing  :: StableSharingTerm -> SharingTerm binder t -> SharingTerm binder t
  TermSharing :: Typeable t
              => StableTermHeight t -> Term binder t       -> SharingTerm binder t

data StableSharingTerm where
  StableSharingTerm :: Typeable t => StableTermHeight t -> SharingTermFloated t -> StableSharingTerm

instance Show StableSharingTerm where
  show (StableSharingTerm sn _) = show $ hashStableTermHeight sn

instance Eq StableSharingTerm where
  StableSharingTerm sn1 _ == StableSharingTerm sn2 _
    | Just sn1' <- gcast sn1 = sn1' == sn2
    | otherwise              = False

higherSST :: StableSharingTerm -> StableSharingTerm -> Bool
StableSharingTerm sn1 _ `higherSST` StableSharingTerm sn2 _ = sn1 `higherSTH` sn2

matchStableTerm :: Typeable t => StableTermHeight t -> StableSharingTerm -> Bool
matchStableTerm sn1 (StableSharingTerm sn2 _)
  | Just sn1' <- gcast sn1 = sn1' == sn2
  | otherwise              = False

noStableTermName :: StableTermHeight t
noStableTermName = unsafePerformIO $ StableTermHeight <$> makeStableName undefined <*> pure 0


  Sharing recovery

  Lambdas are annotated with the level of their binder .
type SharingTermPruned  t = SharingTerm Level t

  Lambdas are annotated by the ' StableSharingTerm ' identifying their binder .
type SharingTermFloated t = SharingTerm StableSharingTerm t

      There is one caveat : We currently rely on ' Tag ' leaves representing free variables to be
recoverSharing :: Typeable t => HOAS.Term t -> SharingTermFloated t
# NOINLINE recoverSharing #
recoverSharing term
  = let (term', occMap) =
          { (term', occMap) <- makeOccMap 0 term
          ; frozenOccMap    <- freezeOccMap occMap
          ; return (term', frozenOccMap)
          }
    in 
    determineScopes occMap term'

  of shared subterms ( Phase One ) .
  importantly , the dependence of the second phase on /global/ occurence information , we may not
  that the ' Tag ' supplied during the initial traversal is already the one required by the HOAS to
  de Bruijn indices , which replaces them by ' VarSharing ' nodes .   Hence , the tag values only serve
  the purpose of determining the ordering in that initial environment . They are /not/ directly used
makeOccMap :: Typeable t => Level -> HOAS.Term t -> IO (SharingTermPruned t, OccMapHash)
makeOccMap lvl rootTerm
  = do
    { occMap <- newOccMapHashTable
    ; (rootTerm', _) <- traverse lvl occMap rootTerm
    ; return (rootTerm', occMap)
    }
  where
    traverse :: forall t. Typeable t
                => Level -> OccMapHash -> HOAS.Term t -> IO (SharingTermPruned t, Int)
    traverse lvl occMap term
      = mfix $ \ ~(_, height) -> do 
  Compute stable name and enter it into the occurence map
        ; sn                        <- makeStableTerm term
        ; heightIfRepeatedOccurence <- enterOcc occMap (StableTermName sn) height
          
        ; traceLine (HOAS.showTermOp term) $
            case heightIfRepeatedOccurence of
              Just height -> "REPEATED occurence (sn = " ++ show (hashStableName sn) ++ 
                             "; height = " ++ show height ++ ")"
              Nothing     -> "first occurence (sn = " ++ show (hashStableName sn) ++ ")"

  NB : This function can only be used in the case alternatives below ; outside of the
      case we can not discharge the '  t ' constraint .
        ; let reconstruct :: Typeable t
                          => IO (Term        Level t, Int)
                          -> IO (SharingTerm Level t, Int)
              reconstruct newTerm
                = case heightIfRepeatedOccurence of 
                    Just height 
                      -> return (VarSharing (StableTermHeight sn height), height)
                    _ -> do
                         { (term, height) <- newTerm
                         ; return (TermSharing (StableTermHeight sn height) term, height)
                         }

        ; case term of
            HOAS.Con v   -> reconstruct $ return (Con v, 1)
            HOAS.Lam f   -> reconstruct $ do
                            ; (body, h) <- traverse (lvl + 1) occMap (f (HOAS.Tag lvl))
                            ; return (Lam lvl body, h + 1)
                            }
            HOAS.App f a -> reconstruct $ do
                            { (f', h1) <- traverse lvl occMap f
                            ; (a', h2) <- traverse lvl occMap a
                            ; return (App f' a', h1 `max` h2 + 1)
                            }
        }

  Type used to maintain how often each shared subterm , so far , occured during a bottom - up sweep .
    - If one shared term 's ' is itself a subterm of another shared term ' t ' , then 's ' must occur
      * after * ' t ' in the ' NodeCounts ' .
    - A term may have a final occurence count of only 1 iff it is either a free variable ( ' Tag ' )
  smaller towards the end of a ' NodeCounts ' list .   The height of free variables ( ' Tag ' ) is 0 ,
  whereas other leaves have height 1 .   This guarantees that all free variables are at the end
  of the ' NodeCounts ' list .
type NodeCounts = [NodeCount]

data NodeCount = NodeCount StableSharingTerm Int
               deriving Show

noNodeCounts :: NodeCounts
noNodeCounts = []

  Singleton node counts for ' Acc '
nodeCount :: StableSharingTerm -> Int -> NodeCounts
nodeCount sst n = [NodeCount sst n]

  Combine node counts that belong to the same node .
(+++) :: NodeCounts -> NodeCounts -> NodeCounts
us +++ vs = foldr insert us vs
  where
    insert x               []                         = [x]
    insert x@(NodeCount st1 count1) ys@(y@(NodeCount st2 count2) : ys') 
      | st1 == st2          = NodeCount (st1 `pickNoneVar` st2) (count1 + count2) : ys'
      | st1 `higherSST` st2 = x : ys
      | otherwise           = y : insert x ys'

    (StableSharingTerm _ (VarSharing _))  `pickNoneVar`  st2  = st2
    st1                                   `pickNoneVar`  _st2 = st1
    
  Determine whether a ' NodeCount ' is for a ' Tag ' , which represent free variables .
isFreeVar :: NodeCount -> Bool
isFreeVar (NodeCount (StableSharingTerm _ (TermSharing _ (Tag  _))) _) = True
isFreeVar _                                                            = False

  Determine the scopes of all variables representing shared subterms ( Phase Two ) in a bottom - up
  Precondition : there are only ' VarSharing ' and ' TermSharing ' nodes in the argument .
determineScopes :: Typeable t
                => OccMap -> SharingTermPruned t -> SharingTermFloated t
determineScopes occMap rootTerm 
  = let
      (sharingTerm, counts) = scopes rootTerm
    in
    if null counts
    then sharingTerm
    else error $ "determineScopes: unbound shared subtrees: " ++ show counts
  where
    scopes :: forall t. SharingTermPruned t -> (SharingTermFloated t, NodeCounts)
    scopes (LetSharing _ _)
      = error $ "determineScopes: scopesAcc: unexpected 'LetSharing'"
    scopes (VarSharing sn)
      = (VarSharing sn, StableSharingTerm sn (VarSharing sn) `nodeCount` 1)
    scopes (TermSharing sn pterm)
      = case pterm of
          Tag i     -> reconstruct (Tag i) noNodeCounts
          Con v     -> reconstruct (Con v) noNodeCounts
          Lam lvl f -> let
                         (f', count)         = scopes f
                         (stableTag, count') = lookupStableTag lvl count
                       in
                       reconstruct (Lam stableTag f') count'
          App f a   -> let
                         (f', count1) = scopes f
                         (a', count2) = scopes a
                       in 
                       reconstruct (App f' a') (count1 +++ count2)
      where
  Occurence count of the currently processed node
        occCount = let StableTermHeight sn' _ = sn
                   in
                   lookupWithTermName occMap (StableTermName sn')

  * If the current node is being shared ( ' occCount > 1 ' ) , replace it by a ' VarSharing '
    node and float the shared subtree out wrapped in a ' NodeCounts ' value .
    float the ' Tag ' out in a ' NodeCounts ' value .   This is idependent of the number of
  In either case , any completed ' NodeCounts ' are injected as bindings using ' LetSharing '
        reconstruct :: Typeable t 
                    => Term StableSharingTerm t -> NodeCounts -> (SharingTermFloated t, NodeCounts)
        reconstruct newTerm@(Tag _) _subCount
          = let thisCount = StableSharingTerm sn (TermSharing sn newTerm) `nodeCount` 1
            in
            tracePure "FREE" (show thisCount) $
            (VarSharing sn, thisCount)
        reconstruct newTerm subCount
          | occCount > 1
          = let allCount = (StableSharingTerm sn sharingTerm `nodeCount` 1) +++ newCount
            in
            tracePure ("SHARED" ++ completed) (show allCount) $
            (VarSharing sn, allCount)
          | otherwise
          = tracePure ("Normal" ++ completed) (show newCount) $
            (sharingTerm, newCount)
          where
            (newCount, bindHere) = filterCompleted subCount
    
  ... and wrap them in ' LetSharing ' constructors
            lets        = foldl (flip (.)) id . map LetSharing $ bindHere
            sharingTerm = lets $ TermSharing sn newTerm
    
            completed | null bindHere = ""
                      | otherwise     = "(" ++ show (length bindHere) ++ " lets)"
    
  NB : Only extract leading nodes ( i.e. , the longest run at the * front * of the list that is
        filterCompleted :: NodeCounts -> (NodeCounts, [StableSharingTerm])
        filterCompleted counts
          = let (completed, counts') = break notComplete counts
            in (counts', [sa | NodeCount sa _ <- completed])
          where
            notComplete nc@(NodeCount st n) | not . isFreeVar $ nc = lookupWithSharingTerm occMap st > n
            notComplete _                                          = True

  continue ( c.f . , comments at ' determineScopes ' ) .
lookupStableTag :: Level -> NodeCounts -> (StableSharingTerm, NodeCounts)
lookupStableTag lvl counts 
  = case partition hasLevel counts of
      (counts', _)    -> 
        error $ "lookupStableTag: duplicate 'Tag's\n  " ++ 
                intercalate ", " [showST st | NodeCount st _ <- counts']
  where
    hasLevel (NodeCount (StableSharingTerm _ (TermSharing _ (Tag lvl'))) _) = lvl == lvl'
    hasLevel (NodeCount st _)                                               = False
        
    noStableSharing :: StableSharingTerm
    noStableSharing = StableSharingTerm noStableTermName (undefined :: SharingTermFloated ())

    showST (StableSharingTerm _ (TermSharing  sn term)) = show (hashStableTermHeight sn) ++ ": " ++ 
                                                          showTermOp term
    showST (StableSharingTerm _ (VarSharing sn))        = "VarSharing " ++ show (hashStableTermHeight sn)
    showST (StableSharingTerm _ (LetSharing st _ ))     = "LetSharing " ++ show st ++ "..."



traceOn :: Bool
traceOn = False

tracePure :: String -> String -> a -> a
tracePure header msg val | traceOn   = trace (header ++ ": " ++ msg) val
                         | otherwise = val

traceLine :: String -> String -> IO ()
traceLine header msg = tracePure header msg $ return ()
