{-
Copyright        : (c) Galois, Inc 2017
Maintainer       : Joe Hendrix <jhendrix@galois.com>

This module provides a function for folding over the subexpressions in
a value without revisiting shared subterms.
-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Macaw.Fold
  ( Data.Parameterized.TraversableFC.FoldableFC(..)
  , ValueFold(..)
  , emptyValueFold
  , foldValueCached
  ) where

import           Control.Monad.State.Strict
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Parameterized.NatRepr
import           Data.Parameterized.Some
import           Data.Parameterized.TraversableFC

import           Data.Macaw.CFG

data ValueFold arch ids r = ValueFold
  { forall arch ids r.
ValueFold arch ids r -> forall (tp :: Type). CValue arch tp -> r
foldCValue  :: !(forall tp . CValue arch tp -> r)
  , forall arch ids r.
ValueFold arch ids r -> forall (utp :: Type). ArchReg arch utp -> r
foldInput      :: !(forall utp . ArchReg arch utp -> r)
  , forall arch ids r.
ValueFold arch ids r
-> forall (utp :: Type). AssignId ids utp -> r -> r
foldAssign     :: !(forall utp . AssignId ids utp -> r -> r)
  }

-- | Empty value fold returns mempty for each non-recursive fold, and the
-- identify of @foldAssign@
emptyValueFold :: Monoid r => ValueFold arch ids r
emptyValueFold :: forall r arch ids. Monoid r => ValueFold arch ids r
emptyValueFold =
  ValueFold { foldCValue :: forall (tp :: Type). CValue arch tp -> r
foldCValue = \CValue arch tp
_ -> r
forall a. Monoid a => a
mempty
            , foldInput :: forall (utp :: Type). ArchReg arch utp -> r
foldInput  = \ArchReg arch utp
_ -> r
forall a. Monoid a => a
mempty
            , foldAssign :: forall (utp :: Type). AssignId ids utp -> r -> r
foldAssign = \AssignId ids utp
_ r
r -> r
r
            }

-- | This folds over elements of a values in a  values.
--
-- It memoizes the results so that if an assignment is visited
-- multiple times, we only visit the children the first time it is
-- visited.  On subsequent visits, `foldAssign` will still be called,
-- but the children will not be revisited.
--
-- This makes the total time to visit linear with respect to the
-- number of children, but still allows determining whether a term is
-- shared.
foldValueCached :: forall r arch ids tp
                .  (Monoid r, FoldableFC (ArchFn arch))
                => ValueFold arch ids r
                -> Value arch ids tp
                -> State (Map (Some (AssignId ids)) r) r
foldValueCached :: forall r arch ids (tp :: Type).
(Monoid r, FoldableFC (ArchFn arch)) =>
ValueFold arch ids r
-> Value arch ids tp -> State (Map (Some (AssignId ids)) r) r
foldValueCached ValueFold arch ids r
fns = Value arch ids tp -> State (Map (Some (AssignId ids)) r) r
forall (tp' :: Type).
Value arch ids tp' -> State (Map (Some (AssignId ids)) r) r
go
  where
    go :: forall tp'
       .  Value arch ids tp'
       -> State (Map (Some (AssignId ids)) r) r
    go :: forall (tp' :: Type).
Value arch ids tp' -> State (Map (Some (AssignId ids)) r) r
go Value arch ids tp'
v =
      case Value arch ids tp'
v of
        CValue CValue arch tp'
c  ->
          r -> State (Map (Some (AssignId ids)) r) r
forall a. a -> StateT (Map (Some (AssignId ids)) r) Identity a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (r -> State (Map (Some (AssignId ids)) r) r)
-> r -> State (Map (Some (AssignId ids)) r) r
forall a b. (a -> b) -> a -> b
$! ValueFold arch ids r -> forall (tp :: Type). CValue arch tp -> r
forall arch ids r.
ValueFold arch ids r -> forall (tp :: Type). CValue arch tp -> r
foldCValue ValueFold arch ids r
fns CValue arch tp'
c
        AssignedValue (Assignment AssignId ids tp'
a AssignRhs arch (Value arch ids) tp'
rhs) -> do
          Map (Some (AssignId ids)) r
m <- StateT
  (Map (Some (AssignId ids)) r)
  Identity
  (Map (Some (AssignId ids)) r)
forall s (m :: Type -> Type). MonadState s m => m s
get
          case Some (AssignId ids) -> Map (Some (AssignId ids)) r -> Maybe r
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (AssignId ids tp' -> Some (AssignId ids)
forall k (f :: k -> Type) (x :: k). f x -> Some f
Some AssignId ids tp'
a) Map (Some (AssignId ids)) r
m of
            Just r
v' ->
              r -> State (Map (Some (AssignId ids)) r) r
forall a. a -> StateT (Map (Some (AssignId ids)) r) Identity a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (r -> State (Map (Some (AssignId ids)) r) r)
-> r -> State (Map (Some (AssignId ids)) r) r
forall a b. (a -> b) -> a -> b
$! ValueFold arch ids r
-> forall (utp :: Type). AssignId ids utp -> r -> r
forall arch ids r.
ValueFold arch ids r
-> forall (utp :: Type). AssignId ids utp -> r -> r
foldAssign ValueFold arch ids r
fns AssignId ids tp'
a r
v'
            Maybe r
Nothing -> do
              r
rhsVal <- (forall (x :: Type).
 r -> Value arch ids x -> State (Map (Some (AssignId ids)) r) r)
-> r
-> AssignRhs arch (Value arch ids) tp'
-> State (Map (Some (AssignId ids)) r) r
forall {k} {l} (t :: (k -> Type) -> l -> Type) (m :: Type -> Type)
       b (f :: k -> Type) (c :: l).
(FoldableFC t, Monad m) =>
(forall (x :: k). b -> f x -> m b) -> b -> t f c -> m b
foldlMFC' (\r
s Value arch ids x
v' -> r -> r -> r
forall a. Monoid a => a -> a -> a
mappend r
s (r -> r)
-> State (Map (Some (AssignId ids)) r) r
-> State (Map (Some (AssignId ids)) r) r
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Value arch ids x -> State (Map (Some (AssignId ids)) r) r
forall (tp' :: Type).
Value arch ids tp' -> State (Map (Some (AssignId ids)) r) r
go Value arch ids x
v') r
forall a. Monoid a => a
mempty AssignRhs arch (Value arch ids) tp'
rhs
              (Map (Some (AssignId ids)) r -> Map (Some (AssignId ids)) r)
-> StateT (Map (Some (AssignId ids)) r) Identity ()
forall s (m :: Type -> Type). MonadState s m => (s -> s) -> m ()
modify' ((Map (Some (AssignId ids)) r -> Map (Some (AssignId ids)) r)
 -> StateT (Map (Some (AssignId ids)) r) Identity ())
-> (Map (Some (AssignId ids)) r -> Map (Some (AssignId ids)) r)
-> StateT (Map (Some (AssignId ids)) r) Identity ()
forall a b. (a -> b) -> a -> b
$ Some (AssignId ids)
-> r -> Map (Some (AssignId ids)) r -> Map (Some (AssignId ids)) r
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (AssignId ids tp' -> Some (AssignId ids)
forall k (f :: k -> Type) (x :: k). f x -> Some f
Some AssignId ids tp'
a) r
rhsVal
              r -> State (Map (Some (AssignId ids)) r) r
forall a. a -> StateT (Map (Some (AssignId ids)) r) Identity a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (r -> State (Map (Some (AssignId ids)) r) r)
-> r -> State (Map (Some (AssignId ids)) r) r
forall a b. (a -> b) -> a -> b
$! ValueFold arch ids r
-> forall (utp :: Type). AssignId ids utp -> r -> r
forall arch ids r.
ValueFold arch ids r
-> forall (utp :: Type). AssignId ids utp -> r -> r
foldAssign ValueFold arch ids r
fns AssignId ids tp'
a r
rhsVal
        Initial ArchReg arch tp'
r    ->
          r -> State (Map (Some (AssignId ids)) r) r
forall a. a -> StateT (Map (Some (AssignId ids)) r) Identity a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (r -> State (Map (Some (AssignId ids)) r) r)
-> r -> State (Map (Some (AssignId ids)) r) r
forall a b. (a -> b) -> a -> b
$! ValueFold arch ids r -> forall (utp :: Type). ArchReg arch utp -> r
forall arch ids r.
ValueFold arch ids r -> forall (utp :: Type). ArchReg arch utp -> r
foldInput ValueFold arch ids r
fns ArchReg arch tp'
r