packages feed

lhc-0.10: src/Grin/HPT/QuickSolve.hs

{-# LANGUAGE OverloadedStrings #-}
module Grin.HPT.QuickSolve
    ( solve
    ) where

import Grin.Types                         ( Renamed(..), NodeType(..) )

import qualified Data.Map as Map
import qualified Data.Set as Set
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Monad.State.Strict

import Grin.HPT.Environment as Env
import Grin.HPT.Interface as Interface
import qualified Grin.HPT.Interface as Interface

import Grin.Stage2.Pretty (ppRenamed)

--import Tick
import Debug.Trace

import Control.Parallel.Strategies

type M a = State HeapAnalysis a

type SharingMap = Map.Map Lhs Bool


solve :: Equations -> ([HeapAnalysis], HeapAnalysis)
solve eqs
    = let eqPairs = Map.toList eqs
          iterate i ls
              = forM_ ls $ \(lhs,rhs) ->
                  do debugMsg $ "Reducing: " ++ ppLhs lhs ++ " " ++ show i
                     reducedRhs <- reduceEqs rhs
                     addReduced lhs reducedRhs
          loop iter prev
              = case execState (debugMsg ("Iteration: " ++ show iter) >> iterate iter eqPairs) prev of
                  (newData) -> -- | rnf newData `seq` True ->
                    let (iterList, finishedData) = if prev == newData then ([newData], newData) else loop (iter+1) newData
                    in (newData : iterList, finishedData)
      in loop 1 (mkHeapAnalysis (Map.map (const mempty) eqs) (nonlinearVariables eqs))

-- Scan for shared variables. A variable is shared if it is used more than once.
-- Detecting shared heap points is done later when we solve the equations.
nonlinearVariables :: Equations -> SharingMap
nonlinearVariables eqs
    = appEndo (execWriter (mapM_ rhsFn (Map.elems eqs))) Map.empty
    where rhsFn (Rhs values) = mapM_ worker values
          pushIdent ident = tell $ Endo $ Map.insertWith (\_ _ -> True) (VarEntry ident) False
          worker (Extract ident (tag, _nt, _missing) _nth)   = pushIdent ident >> pushIdent tag
          worker (ExtractVector ident _nth) = pushIdent ident
          worker (Eval ident)               = pushIdent ident
          worker (Update a b)               = pushIdent a >> pushIdent b
          worker (Apply a b)                = pushIdent a >> pushIdent b
          worker (PartialApply a b)         = return ()
          worker (Ident ident)              = pushIdent ident
          worker (Fetch ident)              = pushIdent ident
          worker Env.Base                   = return ()
          worker Env.Heap{}                 = return ()
          worker (Tag tag _nt _nargs args)  = pushIdent tag >> mapM_ rhsFn args
          worker (VectorTag args)           = mapM_ rhsFn args

debugMsg :: String -> M ()
debugMsg str
    = return () -- trace str (return ())

ppLhs :: Lhs -> String
ppLhs (VarEntry v)   = show (ppRenamed v)
ppLhs (HeapEntry hp) = "@" ++ show hp


addReduced :: Lhs -> Interface.Rhs -> M ()
addReduced lhs rhs
    = do orig <- {-addTick "AddReduced" $ -} lookupEq lhs
         let noNewChanges = rhs `Interface.isSubsetOf` orig
         unless noNewChanges $
           do {-addTick "HPT: Change" $ -}
              modify $ \hpt -> hptAddBinding lhs rhs hpt
              debugMsg $ ppLhs lhs ++ ":"
              --debugMsg $ "Old: " ++ show orig
              --debugMsg $ "Rhs: " ++ show rhs
              debugMsg $ "New: " ++ show (mappend orig rhs)
              shared <- isShared lhs
              when shared $
                mapM_ setShared (listHeapPointers rhs)

listHeapPointers :: Interface.Rhs -> [HeapPointer]
listHeapPointers (Interface.Heap hps) = Set.toList hps
listHeapPointers _ = []


reduceEqs :: Env.Rhs -> M Interface.Rhs
reduceEqs (Rhs rhs) = do rhs' <- mapM reduceEq rhs
                         return $ mconcat rhs'

reduceEq :: RhsValue -> M Interface.Rhs
reduceEq Env.Base  = return $ Interface.Base
reduceEq (Env.Heap hp) = return $ Interface.Heap (Set.singleton hp)
reduceEq (Ident i) = lookupEq (VarEntry i)
reduceEq (Extract eq node n) = reduceExtract eq node n
reduceEq (ExtractVector eq n)
    = do rhs <- lookupEq (VarEntry eq)
         case rhs of
           Interface.Empty -> return mempty
           Interface.Other {rhsVector = args} ->
             return (args `nth` n)
    where nth [] n = error $ "reduceEq: ExtractVector: " ++ show (eq, n)
          nth (x:xs) 0 = x
          nth (x:xs) n = nth xs (n-1)
reduceEq (Tag t nt missing args)
    = do args' <- mapM reduceEqs args
         return $ Other (Map.singleton (t, nt, missing) args') []
reduceEq (VectorTag args)
    = do args' <- mapM reduceEqs args
         return $ Interface.Other Map.empty args'
reduceEq (Eval i) = reduceEval i
reduceEq (Fetch i)
    = do hpt <- get
         return $ lookupHeap i hpt
reduceEq (Apply a b) = reduceApply a b
reduceEq (PartialApply a b)
    = do rhs <- lookupEq (VarEntry a)
         case rhs of
           Empty -> return Empty
           Other{rhsTagged = nodes} ->
             do let f ((tag, nt, n), args)
                      | n == 0    = return mempty
                      | otherwise = do bRhs <- lookupEq (VarEntry b)
                                       return $ Other (Map.singleton (tag, nt, (n-1)) (args ++ [bRhs])) []
                    f t             = error $ "reduceEq: apply: " ++ show t
                liftM mconcat $ mapM f (Map.toList nodes)
reduceEq (Update hp val)
    = do rhs <- lookupEq (VarEntry hp)
         case rhs of
           Interface.Empty -> return mempty
           Interface.Heap hps ->
             do valRhs  <- lookupEq (VarEntry val)
                forM_ (Set.toList hps) $ \hp -> addReduced (HeapEntry hp) valRhs
                return mempty

reduceExtract eq node n
    = do rhs <- lookupEq (VarEntry eq)
         case rhs of
           Interface.Empty -> return mempty
           Other{rhsTagged = nodes} ->
             return (Map.findWithDefault [] node nodes `nth` n)
    where nth [] n = mempty
          nth (x:xs) 0 = x
          nth (x:xs) n = nth xs (n-1)

reduceEval i
    = do hpt <- get
         case lookupLhs (VarEntry i) hpt of
           Interface.Base -> return Interface.Base
           Interface.Empty -> return Interface.Empty
           Interface.Heap hps ->
             do let anyShared = heapIsShared i hpt
                let fn hp = do let worker ((t, FunctionNode, 0), args) = do rhs <- lookupEq (VarEntry t)
                                                                            when (anyShared && rhs /= mempty) $
                                                                              addReduced (HeapEntry hp) rhs
                                                                            return rhs
                                   worker ((t, nt, missing), args)     = return $ Other (Map.singleton (t, nt, missing) args) []
                               case lookupLhs (HeapEntry hp) hpt of
                                 Empty        -> return mempty
                                 Other{rhsTagged = nodes} -> liftM mconcat $ mapM worker (Map.toList nodes)
                liftM mconcat $ mapM fn (Set.toList hps)
           rhs -> error $ "Eval: " ++ show (rhs, i)

reduceApply a b
    = do rhs <- lookupEq (VarEntry a)
         case rhs of
          Empty -> return Empty
          Other{rhsTagged = nodes} ->
            do let f ((func, FunctionNode, 1), args)
                     = reduceEq (Ident func)
                   f ((conc, nt, n), args)
                       | n == 0    = return mempty
                       | otherwise = do bRhs <- lookupEq (VarEntry b)
                                        return $ Other (Map.singleton (conc, nt, (n-1)) (args ++ [bRhs])) []
               liftM mconcat $ mapM f (Map.toList nodes)



-- FIXME: Throw an exception if 'lhs' couldn't be found.
lookupEq :: Lhs -> M Interface.Rhs
lookupEq lhs
    = gets $ \(hpt) -> lookupLhs lhs hpt

-- FIXME: Throw an exception if 'lhs' couldn't be found.
isShared :: Lhs -> M Bool
isShared lhs
    = gets $ \(hpt) -> hptIsShared lhs hpt

setShared :: HeapPointer -> M ()
setShared hp = modify $ \hpt -> hptSetShared (HeapEntry hp) hpt

{-
lhsIsDead :: Lhs -> M Bool
lhsIsDead lhs
    = asks $ \(_hpt, dead) -> lhs `Set.member` dead

lhsSetDead :: Lhs -> M ()
lhsSetDead lhs
    = tell (mempty, Endo $ Set.insert lhs)
-}