<commit_before>{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE DeriveDataTypeable   #-}
{-# LANGUAGE DeriveFunctor        #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE ImplicitParams       #-}
{-# LANGUAGE NamedFieldPuns       #-}
{-# LANGUAGE Rank2Types           #-}
{-# LANGUAGE RecursiveDo          #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances #-}


-- This module implements a hybrid representation of dags, that uses
-- explicit edges only if necessary (i.e. in case of sharing). 

module Dag where

import Control.Applicative
import Control.Exception.Base
import Control.Monad.ST
import Control.Monad.State
import qualified Data.HashMap.Lazy as HashMap
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IORef
import Data.Maybe
import Data.STRef
import Data.Traversable (Traversable)
import qualified Data.Traversable as Traversable
import Data.Typeable
import Data.Vector (MVector, Vector)
import qualified Data.Vector as Vec
import qualified Data.Vector.Generic.Mutable as MVec
import System.Mem.StableName

import AG


-- | The type of node in a 'Dag'.

type Node = Int

-- | The type of the compact edge representation used in a 'Dag'.

type Edges f = IntMap (f (Free f Node))

-- | The type of directed acyclic graphs (DAGs). 'Dag's are used as a
-- compact representation of 'Tree's.

data Dag f = Dag
    { root      :: f (Free f Node) -- ^ the entry point for the DAG
    , edges     :: Edges f         -- ^ the edges of the DAG
    , nodeCount :: Int             -- ^ the total number of nodes in the DAG
    }

-- | Construct a Dag
mkDag :: f (Free f Node) -> [(Node, f (Free f Node))] -> Dag f
mkDag r es = Dag r (IntMap.fromList es) (if null es then 0 else maximum (map fst es) + 1)

-- | This function unravels a given graph to the term it
-- represents.

unravelDag :: forall f. Functor f => Dag f -> Tree f
unravelDag Dag {edges, root} = In $ build <$> root
    where build :: Free f Node -> Tree f
          build (In t) = In $ build <$> t
          build (Ret n) = In $ build <$> edges IntMap.! n

-- | This exception indicates that a 'Tree' could not be reified to a
-- 'Dag' (using 'reifyDag') due to its cyclic sharing structure.
data CyclicException = CyclicException
    deriving (Show, Typeable)

instance Exception CyclicException

-- | This function takes a term, and returns a 'Dag' with the implicit
-- sharing of the input data structure made explicit. If the sharing
-- structure of the term is cyclic an exception of type
-- 'CyclicException' is thrown.
reifyDag :: Traversable f => Tree f -> IO (Dag f)
reifyDag m = do
  tabRef <- newIORef HashMap.empty
  let findNodes (In !j) = do
        st <- liftIO $ makeStableName j
        tab <- readIORef tabRef
        case HashMap.lookup st tab of
          Just (single,f) | single -> writeIORef tabRef (HashMap.insert st (False,f) tab)
                                      >> return st
                          | otherwise -> return st
          Nothing -> do res <- Traversable.mapM findNodes j
                        tab <- readIORef tabRef
                        if HashMap.member st tab
                          then throwIO CyclicException
                          else writeIORef tabRef (HashMap.insert st (True,res) tab)
                               >> return st
  st <- findNodes m
  tab <- readIORef tabRef
  counterRef <- newIORef 0
  edgesRef <- newIORef IntMap.empty
  nodesRef <- newIORef HashMap.empty
  let run st = do
        let (single,f) = tab HashMap.! st
        if single then In <$> Traversable.mapM run f
        else do
          nodes <- readIORef nodesRef
          case HashMap.lookup st nodes of
            Just n -> return (Ret n)
            Nothing -> do
              n <- readIORef counterRef
              writeIORef counterRef $! (n+1)
              writeIORef nodesRef (HashMap.insert st n nodes)
              f' <- Traversable.mapM run f
              modifyIORef edgesRef (IntMap.insert n f')
              return (Ret n)
  In root <- run st
  edges <- readIORef edgesRef
  count <- readIORef counterRef
  return (Dag root edges count)


-- | This function runs an AG on a dag.
runAGDag :: forall f d u .Traversable f
    => (d -> d -> d)   -- ^ Resolution of top-down state
    -> Syn' f (u,d) u  -- ^ Bottom-up state propagation
    -> Inh' f (u,d) d  -- ^ Top-down state propagation
    -> (u -> d)         -- ^ Initial top-down state
    -> Dag f
    -> u
runAGDag res syn inh dinit g = uFin where
    uFin = umapFin IntMap.! root g
    dFin = dinit uFin
    run :: d -> Free f Node -> (u, IntMap d)
    run d (Ret x) = (umapFin IntMap.! x, IntMap.singleton x d)
    run d (In t)  = (u, dmapLoc) where
        u = explicit syn (u,d) unNumbered result
        m = explicit inh (u,d) unNumbered result
        (result, (dmapLoc,_)) = runState (Traversable.mapM run' t) (IntMap.empty,0)
        run' :: Free f Node -> State (IntMap d, Int) (Numbered ((u,d)))
        run' s = do
            (oldDmap,i) <- get
            let d' = lookupNumMap d i m
                (u',dmap') = run d' s
            let j = i+1
            j `seq` put (IntMap.unionWith res dmap' oldDmap, j)
            return (Numbered i (u',d'))
    dmapFin = IntMap.foldr (\ (_,m1) m2 -> IntMap.unionWith res m1 m2) 
           (IntMap.singleton (root g) dFin) result
    umapFin = IntMap.map fst result
    result = IntMap.mapWithKey (\ n t -> run (dmapFin IntMap.! n) t) (edges g)


-- | This function runs an attribute grammar on a dag. The result is
-- the (combined) synthesised attribute at the root of the dag.

runAGDagST :: forall f d u .Traversable f
    => (d -> d -> d)   -- ^ resolution function for inherited attributes
    -> Syn' f (u,d) u  -- ^ semantic function of synthesised attributes
    -> Inh' f (u,d) d  -- ^ semantic function of inherited attributes
    -> (u -> d)        -- ^ initialisation of inherited attributes
    -> Dag f           -- ^ input dag
    -> u
runAGDagST _ syn inh dinit Dag {root = t,nodeCount=0} = uFin where
    dFin = dinit uFin
    uFin = run dFin t 
    run d t = u where
        t' = fmap bel $ number t
        bel (Numbered i s) =
            let d' = lookupNumMap d i m
            in Numbered i (run' d' s, d')
        m = explicit inh (u,d) unNumbered t'
        u = explicit syn (u,d) unNumbered t'
    run' _ (Ret _) = error "runAGDagST: the input dag has no edges!"
    run' d (In t) = run d t<commit_msg>Revert "refactor implementation of runAGDag"<commit_after>{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE DeriveDataTypeable   #-}
{-# LANGUAGE DeriveFunctor        #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE ImplicitParams       #-}
{-# LANGUAGE NamedFieldPuns       #-}
{-# LANGUAGE Rank2Types           #-}
{-# LANGUAGE RecursiveDo          #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances #-}


-- This module implements a hybrid representation of dags, that uses
-- explicit edges only if necessary (i.e. in case of sharing). 

module Dag where

import Control.Applicative
import Control.Exception.Base
import Control.Monad.ST
import Control.Monad.State
import qualified Data.HashMap.Lazy as HashMap
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IORef
import Data.Maybe
import Data.STRef
import Data.Traversable (Traversable)
import qualified Data.Traversable as Traversable
import Data.Typeable
import Data.Vector (MVector, Vector)
import qualified Data.Vector as Vec
import qualified Data.Vector.Generic.Mutable as MVec
import System.Mem.StableName

import AG


-- | The type of node in a 'Dag'.

type Node = Int

-- | The type of the compact edge representation used in a 'Dag'.

type Edges f = IntMap (f (Free f Node))

-- | The type of directed acyclic graphs (DAGs). 'Dag's are used as a
-- compact representation of 'Tree's.

data Dag f = Dag
    { root      :: f (Free f Node) -- ^ the entry point for the DAG
    , edges     :: Edges f         -- ^ the edges of the DAG
    , nodeCount :: Int             -- ^ the total number of nodes in the DAG
    }

-- | Construct a Dag
mkDag :: f (Free f Node) -> [(Node, f (Free f Node))] -> Dag f
mkDag r es = Dag r (IntMap.fromList es) (if null es then 0 else maximum (map fst es) + 1)

-- | This function unravels a given graph to the term it
-- represents.

unravelDag :: forall f. Functor f => Dag f -> Tree f
unravelDag Dag {edges, root} = In $ build <$> root
    where build :: Free f Node -> Tree f
          build (In t) = In $ build <$> t
          build (Ret n) = In $ build <$> edges IntMap.! n

-- | This exception indicates that a 'Tree' could not be reified to a
-- 'Dag' (using 'reifyDag') due to its cyclic sharing structure.
data CyclicException = CyclicException
    deriving (Show, Typeable)

instance Exception CyclicException

-- | This function takes a term, and returns a 'Dag' with the implicit
-- sharing of the input data structure made explicit. If the sharing
-- structure of the term is cyclic an exception of type
-- 'CyclicException' is thrown.
reifyDag :: Traversable f => Tree f -> IO (Dag f)
reifyDag m = do
  tabRef <- newIORef HashMap.empty
  let findNodes (In !j) = do
        st <- liftIO $ makeStableName j
        tab <- readIORef tabRef
        case HashMap.lookup st tab of
          Just (single,f) | single -> writeIORef tabRef (HashMap.insert st (False,f) tab)
                                      >> return st
                          | otherwise -> return st
          Nothing -> do res <- Traversable.mapM findNodes j
                        tab <- readIORef tabRef
                        if HashMap.member st tab
                          then throwIO CyclicException
                          else writeIORef tabRef (HashMap.insert st (True,res) tab)
                               >> return st
  st <- findNodes m
  tab <- readIORef tabRef
  counterRef <- newIORef 0
  edgesRef <- newIORef IntMap.empty
  nodesRef <- newIORef HashMap.empty
  let run st = do
        let (single,f) = tab HashMap.! st
        if single then In <$> Traversable.mapM run f
        else do
          nodes <- readIORef nodesRef
          case HashMap.lookup st nodes of
            Just n -> return (Ret n)
            Nothing -> do
              n <- readIORef counterRef
              writeIORef counterRef $! (n+1)
              writeIORef nodesRef (HashMap.insert st n nodes)
              f' <- Traversable.mapM run f
              modifyIORef edgesRef (IntMap.insert n f')
              return (Ret n)
  In root <- run st
  edges <- readIORef edgesRef
  count <- readIORef counterRef
  return (Dag root edges count)


-- | This function runs an AG on a dag.
runAGDag :: forall f d u .Traversable f
    => (d -> d -> d)   -- ^ Resolution of top-down state
    -> Syn' f (u,d) u  -- ^ Bottom-up state propagation
    -> Inh' f (u,d) d  -- ^ Top-down state propagation
    -> (u -> d)         -- ^ Initial top-down state
    -> Dag f
    -> u
runAGDag res syn inh dinit g = u
    where d = dinit u
          syn' :: SynExpl f (u,d) u
          syn' = explicit syn
          inh' :: InhExpl f (u,d) d
          inh' = explicit inh
          run = free res  syn' inh'
          (u,dmapRoot) = run d umap (root g)
          dmap = IntMap.foldr (\ (_,m1) m2 -> IntMap.unionWith res m1 m2)
                 dmapRoot result
          umap = IntMap.map fst result
          result = IntMap.mapWithKey (\ n t -> run (dmap IntMap.! n) umap t) (edges g)


-- | Auxiliary function for 'runAGDag'.

free :: forall f u d . Traversable f => (d -> d -> d) -> SynExpl f (u,d) u -> InhExpl f (u,d) d
     -> d -> IntMap u -> f (Free f Node) -> (u, IntMap d)
free res syn inh d umap s = run d s where
    run :: d -> f (Free f Node) -> (u, IntMap d)
    run d t = (u, dmap)
        where u = syn (u,d) unNumbered result
              m = inh (u,d) unNumbered result
              (result, (dmap,_)) = runState (Traversable.mapM run' t) (IntMap.empty,0)
              run' :: Free f Node -> State (IntMap d, Int) (Numbered ((u,d)))
              run' s = do
                  (oldDmap,i) <- get
                  let d' = lookupNumMap d i m
                      (u',dmap') = runF d' s
                  put (IntMap.unionWith res dmap' oldDmap, (i+1))
                  return (Numbered i (u',d'))
    runF :: d -> Free f Node -> (u, IntMap d)
    runF d (Ret x) = (umap IntMap.! x, IntMap.singleton x d)
    runF d (In t)  = run d t

-- | This function runs an attribute grammar on a dag. The result is
-- the (combined) synthesised attribute at the root of the dag.

runAGDagST :: forall f d u .Traversable f
    => (d -> d -> d)   -- ^ resolution function for inherited attributes
    -> Syn' f (u,d) u  -- ^ semantic function of synthesised attributes
    -> Inh' f (u,d) d  -- ^ semantic function of inherited attributes
    -> (u -> d)        -- ^ initialisation of inherited attributes
    -> Dag f           -- ^ input dag
    -> u
runAGDagST _ syn inh dinit Dag {root = t,nodeCount=0} = uFin where
    dFin = dinit uFin
    uFin = run dFin t 
    run d t = u where
        t' = fmap bel $ number t
        bel (Numbered i s) =
            let d' = lookupNumMap d i m
            in Numbered i (run' d' s, d')
        m = explicit inh (u,d) unNumbered t'
        u = explicit syn (u,d) unNumbered t'
    run' _ (Ret _) = error "runAGDagST: the input dag has no edges!"
    run' d (In t) = run d t
{'fraction_non_alphanumeric': 0.10431893687707641, 'fraction_numerical': 0.0017718715393133999, 'mean_word_length': 2.0419941612396135, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 100, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6547667', 'n_tokens_mistral': 4425, 'n_tokens_neox': 4258, 'n_words': 1664}