content: {-# OPTIONS_GHC -funbox-strict-fields #-}
  
  -- An implementation of a Thyer lazy specializer. From:
  -- Lazy Specialization
  -- by Michael Jonathan Thyer (1999).
  
  -- Memoizing substitutions not implemented.
  
  module Thyer (eval) where
  
  import qualified Depth
  import qualified HOAS
  import qualified IORefRef as Ref
  import Control.Applicative
  import Control.Monad ((<=<))
  
  data Blocked
      = Blocked
      | Unblocked
      deriving (Eq)
  
  type NodeRef a = Ref.Ref (Node a)
  
  data Node a = Node {
      nodeBlocked :: !Blocked,
      nodeDepth   :: !Int,
      nodeData    :: !(NodeData a)
    }
  
  data NodeData a
      = Lambda !(NodeRef a)
      | Apply  !(NodeRef a) !(NodeRef a)
      | Subst  !(NodeRef a) !Int !(NodeRef a) !Int   -- body var arg shift
      | Var
      | Prim   a
  
  -- reduce reduces its argument to whnf *destructively*.  It returns the reduced 
  -- node for convenience.  reduce x = reduce x >> Ref.read x.
  reduce :: (HOAS.Primitive a) => NodeRef a -> IO (Node a)
  reduce ref = do
      node <- Ref.read ref
      if nodeBlocked node == Blocked then return node else do
      case nodeData node of
          Apply f arg -> do
              fnode <- reduce f
              case nodeData fnode of
                  Lambda body -> do
                      let bind = nodeDepth fnode + 1
  
                          -- shift is the amount by which the depths of f's nodes are expected
                          -- to change.
                          shift = nodeDepth node - bind
  
                          -- This is (nodeDepth fnode) instead of (nodeDepth node) in Thyer's
                          -- paper. It is somewhat irrelevant since we only check depths when we
                          -- are substituting through a node, and we never subsitute through a 
                          -- subst node, but I believe this makes more sense.
                          node' = Node Unblocked (nodeDepth node) (Subst body bind arg shift)
                      Ref.write ref node'
                      reduce ref
                  Prim p -> do
                      argnode <- reduce arg
                      case nodeData argnode of
                          Prim p'   -> sideEffect (Ref.write ref) (Node Blocked 0 (Prim $ p `HOAS.apply` p'))
                          Apply {}  -> blocked
                          Var {}    -> blocked
                          Lambda {} -> fail "Can't apply primitive to lambda"
                          Subst {}  -> fail "Bug - reduced expression ended up a subst"
                  _ -> blocked
          Subst body var arg shift -> do
              -- This is the code that has the specializing effect.  We *reduce*
              -- the body, including application nodes, before substituting into it.  
              -- A simple lazy evaluator would just push down the substitution through
              -- any type of node, including applications.  cf. Thyer p. 122.
              reduce body
              Ref.link ref =<< subst body var arg shift
              reduce ref
          _ -> blocked
      where
      blocked = do
          node <- Ref.read ref
          sideEffect (Ref.write ref) $! node { nodeBlocked = Blocked }
  
  sideEffect :: (a -> IO ()) -> a -> IO a
  sideEffect f x = f x >> return x
  
  -- subst returns body with the variable at depth bind substituted for arg.  It
  -- does not modify its arguments.
  subst :: NodeRef a -> Int -> NodeRef a -> Int -> IO (NodeRef a)
  subst body bind arg shift = do  
      bodynode <- Ref.read body
      -- if the depth of the body is less than the depth of the variable we are 
      -- substituting, the variable cannot possibly occur in the body, so just 
      -- dissolve away.
      if nodeDepth bodynode < bind then return body else do
      
      let newdepth = nodeDepth bodynode + shift
      case nodeData bodynode of
          Var | nodeDepth bodynode == bind -> return arg
              | otherwise                  -> Ref.new (Node Blocked newdepth Var)
          Lambda body -> do
              substbody <- Ref.new (Node Unblocked (newdepth+1) (Subst body bind arg shift))
              Ref.new (Node Unblocked newdepth (Lambda substbody))
          Apply f x -> do
              f' <- Ref.new (Node Unblocked newdepth (Subst f bind arg shift)) 
              x' <- Ref.new (Node Unblocked newdepth (Subst x bind arg shift))
              Ref.new (Node Unblocked newdepth (Apply f' x'))
          _ -> return body
  
  fromDepth :: Depth.ExpNode a -> IO (NodeRef a)
  fromDepth (d, n) = case n of
      Depth.Lambda body -> Ref.new . Node Unblocked d . Lambda =<< fromDepth body
      Depth.Apply f x   -> Ref.new =<< Node Unblocked d <$> liftA2 Apply (fromDepth f) (fromDepth x)
      Depth.Var         -> Ref.new (Node Blocked d Var)
      Depth.Prim x      -> Ref.new . Node Blocked d . Prim $ x
  
  getValue :: (HOAS.Primitive a) => NodeRef a -> IO a
  getValue ref = do
      refnode <- reduce ref
      case nodeData refnode of
          Prim x -> return x
          _ -> fail "Not a value"
  
  eval :: (HOAS.Primitive a) => Depth.Depth a -> IO a
  eval = getValue <=< fromDepth . Depth.getDepth
  