Creating a Visitor Pattern with Polymorphic Recursive Modules
<p>(Disclaimer: I am fairly certain that this is not idiomatic in any way. If there is some alternative tree-traversal idiom in OCaml, I'm all ears :) )</p>

<p>I am writing a toy compiler in OCaml, and I would like to have a visitor for my large syntax tree type. I wrote one using classes, but I thought it would be fun to try and implement one using modules/functors. My type hierarchy is massive, so let me illustrate what I'm trying to do.</p>

<p>Consider the following type definitions (making these up on the spot):</p>

<pre><code>type expr = 
    SNum of int
  | SVarRef of string
  | SAdd of expr * expr
  | SDo of stmt list

and stmt = 
    SIf of expr * expr * expr
  | SAssign of string * expr
</code></pre>

<p>Let me briefly illustrate usage. Say, for example, I wanted to collect all of the <code>SVarRef</code>s inside of the program. If I had a mapping visitor (one which visits every node of the tree and does nothing by default), I could do the following (in a perfect world):</p>

<pre><code>module VarCollector : (sig 
  include AST_VISITOR 
  val var_refs : (string list) ref 
end) = struct
  include MapVisitor

  let var_refs = ref []

  let s_var_ref (s : string) =
    var_refs := s::!var_refs
    SVarRef(s)
end

(* Where my_prog is a stmt type *)
let refs = begin 
  let _ = VarCollector.visit_stmt my_prog in
    VarCollector.!var_refs
end
</code></pre>

<p>I should note that the advantage of having functions for each specific variant is that my actual codebase has a type which both has a large amount of variants and does not make sense to fragment. Variant-specific functions allow the avoiding of repeated iteration implementations for the other variants of a type.</p>

<p>Seems simple enough, but here's the catch: there are different types of visitors. <code>MapVisitor</code> returns the original syntax tree, so it has type</p>

<pre><code>sig
  (** Dispatches to variant implementations *)
  val visit_stmt : stmt -&gt; stmt
  val visit_expr : expr -&gt; expr
  (** Variant implementations *)
  val s_num : int -&gt; expr
  val s_var_ref : string -&gt; expr
  val s_add : (expr * expr) -&gt; expr
  val s_do : stmt list -&gt; expr
  val s_if : (expr * expr * expr) -&gt; stmt
  val s_assign : (string * expr) -&gt; stmt
end
</code></pre>

<p>At the same time, one might imagine a folding visitor in which the return type is some <code>t</code> for every function. Attempting to abstract this as much as possible, here is my attempt:</p>

<pre><code>module type AST_DISPATCHER = sig
  type expr_ret
  type stmt_ret
  val visit_expr : expr -&gt; expr_ret
  val visit_stmt : stmt -&gt; stmt_ret
end
(** Concrete type designation goes in AST_VISITOR_IMPL *)
module type AST_VISITOR_IMPL = sig
  type expr_ret
  type stmt_ret

  val s_num : int -&gt; expr_ret
  val s_var_ref : string -&gt; expr_ret
  val s_add : (expr * expr) -&gt; expr_ret
  val s_do : stmt list -&gt; expr_ret

  val s_if : (expr * expr * expr) -&gt; stmt_ret
  val s_assign : (string * expr) -&gt; stmt_ret
end
module type AST_VISITOR = sig
  include AST_VISITOR_IMPL
  include AST_DISPATCHER with type expr_ret := expr_ret
                          and type stmt_ret := stmt_ret
end

(** Dispatcher Implementation *)
module AstDispatcherF(IM : AST_VISITOR_IMPL) : AST_DISPATCHER = struct
  type expr_ret = IM.expr_ret
  type stmt_ret = IM.stmt_ret

  let visit_expr = function
    | SNum(i) -&gt; IM.s_num i
    | SVarRef(s) -&gt; IM.s_var_ref s
    | SAdd(l,r) -&gt; IM.s_add (l,r)
    | SDo(sl) -&gt; IM.s_do sl

  let visit_stmt = function
    | SIf(c,t,f) -&gt; IM.s_if (c,t,f)
    | SAssign(s,e) -&gt; IM.s_assign (s,e)
end
module rec MapVisitor : AST_VISITOR = struct
  type expr_ret = expr
  type stmt_ret = stmt
  module D : (AST_DISPATCHER with type expr_ret := expr_ret
                              and type stmt_ret := stmt_ret)
    = AstDispatcherF(MapVisitor)

  let visit_expr = D.visit_expr
  let visit_stmt = D.visit_stmt

  let s_num i = SNum i
  let s_var_ref s = SVarRef s
  let s_add (l,r) = SAdd(D.visit_expr l, D.visit_expr r)
  let s_do sl = SDo(List.map D.visit_stmt sl)

  let s_if (c,t,f) = SIf(D.visit_expr c, D.visit_expr t, D.visit_expr f)
  let s_assign (s,e) = SAssign(s, D.visit_expr e)
end
</code></pre>

<p>Running this gives me the following error message, however:</p>

<pre><code>Error: Signature Mismatch:
       Values do not match:
         val visit_expr : expr -&gt; expr_ret
       is not included in
         val visit_expr : expr -&gt; expr_ret
</code></pre>

<p>I know this means that I am not expressing the relationship between the types correctly, but I cannot figure out what the fix is in this case.</p>