A more serious online partial evaluator

Admittedly, the illustrative partial evaluator of the previous blog post was quite limited. Consider, for example, the following attempts of partially evaluating different applications of the exponentiation function; the last attempt diverges:

> peval (lib, (Apply "exp" [Const 2, Const 3]))
Const 8

> peval (lib, (Apply "exp" [Var "x", Const 3]))
Binary Times (Var "x") (Binary Times (Var "x") (Binary Times (Var "x")(Const 1)))

> peval (lib, (Apply "exp" [Const 2, Var "n"]))
IfZero (Var "n") (Const 1) (Binary Times (Const 2) (IfZero (Binary Plus (Var "n") (Const (-1))) (Const 1) (Binary Times (Const 2) (IfZero (Binary Plus (Binary Plus (Var "n") (Const (-1))) (Const (-1))) (Const 1) (Binary Times (Const 2) (IfZero (Binary Plus (Binary Plus (Binary Plus (Var "n") (Const (-1))) (Const (-1))) (Const (-1))) (Const 1) (Binary Times (Const 2) (IfZero (Binary Plus (Binary Plus (Binary Plus (Binary Plus (Var "n") (Const (-1))) (Const (-1))) (Const (-1))) (Const (-1))) (Const 1) (Binary Times (Const 2) (IfZero ...

In the first attempt, we partially evaluate a ground term, which is as if we were using the regular interpreter. In the second attempt, we leave the base parametric, but we fix the exponent. In the third attempt, it is the other way around: we fix the base, but we leave the exponent parametric. The divergence for the third attempt is clearly undesirable. Ideally, partial evaluation of a non-ground term t should terminate whenever there is an ground instance t' of t such that regular evaluation terminates for t'. Otherwise, we could never be sure whether we can safely decompose a computation into partial evaluation followed by evaluation of the residual program.

So let us develop a more advanced partial evaluator:

peval :: Prog -> (Expr, FEnv)

The result type has changed. In addition to a transformed main expression, we also expect to obtain specialized function definitions as they are encountered by partial evaluation. That is, function definitions are specialized by statically known argument values. A demo helps:

> peval (lib, (Apply "exp" [Var "x", Const 3]))
( Apply "exp_[(n,3)]" [Var "x"],
  [("exp_[(n,3)]", (["x"], Apply "times" [Var "x",Apply "exp_[(n,2)]" [Var "x"]])),
   ("exp_[(n,2)]", (["x"], Apply "times" [Var "x",Apply "exp_[(n,1)]" [Var "x"]])),
   ("exp_[(n,1)]", (["x"], Apply "times" [Var "x",Apply "exp_[(n,0)]" [Var "x"]])),
   ("exp_[(n,0)]", (["x"], Const 1))])

Thus, specialized function definitions were inferred for all the inductively encountered values for the exponent. Subject to a simple inlining optimization (which is not shown here for brevity but included into the open-source repository), we obtain the familiar expression for "x" to the power 3:

Apply "times" [Var "x",Apply "times" [Var "x",Apply "times" [Var "x",Const 1]]]

The partial evaluator commences as follows:

peval :: Prog -> (Expr, FEnv)
peval (fe,m) = runState (peval' m []) []
 where
  peval' :: Expr -> VEnv -> State FEnv Expr 

That is, the expression-level partial evaluation function starts from the empty variable environment and it uses the state monad to aggregate specialized function definitions, starting also from the empty list. The type of variable environments is the same as in the original interpreter:

type VEnv = [(String,Int)]

The cases for all constructs but function application can be adopted from the simpler partial evaluator--except that we need to convert to monadic style:

  peval' :: Expr -> VEnv -> State FEnv Expr 
  peval' (Const i) ve = return (Const i)
  peval' (Var x) ve
   = return (case lookup x ve of
              Just v -> Const v
              Nothing -> Var x)
  peval' (Binary op e1 e2) ve
   = do
        e1' <- peval' e1 ve
        e2' <- peval' e2 ve
        return (case (e1', e2') of 
                 (Const v1, Const v2) -> Const (op2f op v1 v2)
                 _ -> Binary op e1' e2')
  peval' (IfZero e1 e2 e3) ve
   = do
        e1' <- peval' e1 ve
        case e1' of 
          (Const v1) -> if (v1==0) then r2 else r3
          _ -> r2 >>= \e2' -> r3 >>= \e3' -> return (IfZero e1' e2' e3')
   where
    r2 = peval' e2 ve
    r3 = peval' e3 ve

The deviating and complicated and interesting case is the one for function application. We prepare this development by an informal specification. (Thanks again to William Cook for teaching me on this part.) The following specification relies on two important terms used in partial evaluation. That is, when partial evaluation encounters a function application with a value (say, a constant---as opposed to an expression) for a given argument position, then this position (or the associated variable) is called static; otherwise, it is called dynamic. The specification follows.

1. The applied function is looked up. 2. The arguments of the function application are partially evaluated. 3. The arguments are partitioned into static and dynamic ones. 4. The specialized function name consists of the original function name qualified by the static arguments. 5. The body of the specialized funtion is obtained by partially evaluating the original body in the variable envrironment of the static variables. 6. An application of the specialized function is constructed from the dynamic arguments. This is the result of partial evaluation.

The following function definition clarifies the step of partitioning the partially evaluated arguments (see 3. above). One list of
expressions is partioned into two lists of static versus dynamic arguments qualified by the formal argument names. Static arguments are Const expressions. Dynamic arguments are all other expressions. Thus:

    partition [] [] = ([],[])
    partition (n:ns) (Const i:es)
     = ((n,i):ss,ds) where (ss,ds) = partition ns es
    partition (n:ns) (e:es)
     = (ss,(n,e):ds) where (ss,ds) = partition ns es

The following equation for function applications implements the afore-mentioned steps 1.-6., but additional details
arise. In order to deal with recursion, it is important that the specialized function is already added to the state before its body is
obtained. To this end, an undefined body is preliminarily associated with the new name, which is replaced later; see get, put, get, put. Furher, we need to distunguish function applications that are fully static from those that involve dynamic arguments because fully static function applications should not be specialized; instead, they should be evaluated.

Here are the glory details:

  peval' (Apply n es) ve
   = do
     -- Look up function
     let (ns,e) = fromJust (lookup n fe)

     -- Partially evaluate arguments
     es' <- mapM (flip peval' ve) es

     -- Partition arguments into static and dynamic ones
     let (ss,ds) = partition ns es'

     case (null ss, null ds) of

       -- A constant application
       (True, True) -> peval' e []

       -- A fully static application
       (False, True) -> peval' e ss

       -- A fully dynamic application
       (True, False) -> return (Apply n es')

       -- A mixed static application
       (False, False) -> (do

         -- The name that encodes the static values
         let n' = n ++ "_" ++ myshowl ss

         -- Construct application of specialized function
         let e' = Apply n' (map snd ds)
        
         -- See whether the specialization exists already           
         fe <- get
         case lookup n' fe of
           Just _ -> return e'
           Nothing -> (do 

             -- Memo before possible recursion
             put (fe++[(n',undefined)])

             -- Partial evaluation of function body
             e'' <- peval' e ss

             -- Record definition of specialized function
             fe' <- get
             put (update (const (map fst ds,e'')) n' fe')

             -- Return application of specialized function
             return e'))

So much for partial evaluation.

Regards,
PF

PS: The source code of this blog post is available through the repository of the Software Language Processing Suite; see here.