Creating a Visitor Pattern with Polymorphic Recursive Modules

Question

(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 :) )

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.

Consider the following type definitions (making these up on the spot):

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

Let me briefly illustrate usage. Say, for example, I wanted to collect all of the SVarRefs 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):

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

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.

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

sig
  (** Dispatches to variant implementations *)
  val visit_stmt : stmt -> stmt
  val visit_expr : expr -> expr
  (** Variant implementations *)
  val s_num : int -> expr
  val s_var_ref : string -> expr
  val s_add : (expr * expr) -> expr
  val s_do : stmt list -> expr
  val s_if : (expr * expr * expr) -> stmt
  val s_assign : (string * expr) -> stmt
end

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

module type AST_DISPATCHER = sig
  type expr_ret
  type stmt_ret
  val visit_expr : expr -> expr_ret
  val visit_stmt : stmt -> 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 -> expr_ret
  val s_var_ref : string -> expr_ret
  val s_add : (expr * expr) -> expr_ret
  val s_do : stmt list -> expr_ret

  val s_if : (expr * expr * expr) -> stmt_ret
  val s_assign : (string * expr) -> 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) -> IM.s_num i
    | SVarRef(s) -> IM.s_var_ref s
    | SAdd(l,r) -> IM.s_add (l,r)
    | SDo(sl) -> IM.s_do sl

  let visit_stmt = function
    | SIf(c,t,f) -> IM.s_if (c,t,f)
    | SAssign(s,e) -> 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

Running this gives me the following error message, however:

Error: Signature Mismatch:
       Values do not match:
         val visit_expr : expr -> expr_ret
       is not included in
         val visit_expr : expr -> expr_ret

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.

Answer 1

Disclaimer: Modules are just records of values accompanied with type definitions. Since there are no types in your modules there is no need to use them at all, just use plain old record types, and you will get one of the idiomatic AST traversal pattern. Soon, you will find out that you need an open recursion, and will switch to class-based approach. Anyway, this was the main reason, why classes were added to OCaml. You may also wrap your classes into a state monad, to fold AST with arbitrary user data.

What concerning your error message, then it is simple, you hide your types with signatures, a common mistake. The easiest solution is to omit an annotation of the return type of a functor at all, or to propagate type equality with a with type expr = expr annotations.

If you need examples of more idiomatic approaches, then for records you can go for ppx mappers, here is an example of different visitors implemented with classes, including those that are wrapped into a state monad.