ppx_hardcaml, using expression mapper

ppx_hardcaml.ml

(*

  build with:
 
 ocamlbuild -use-ocamlfind -package compiler-libs.common -package ppx_tools.metaquot ppx_hardcaml.native 
 
*)

open Ast_mapper
open Ast_helper
open Ast_convenience
open Asttypes
open Parsetree
open Longident

open Printf

let location_exn ~loc msg =
  Location.Error (Location.error ~loc msg)
  |> raise

let rec hwexpr mapper expr =
  (* TODO resize operators ... *)
  (*let uresize op a b = 
    let op = {pexp_desc = Pexp_ident {txt = Lident op; loc=default_loc}} in
    [%expr 
      let a,b = [%e a], [%e b] in
      let w = max (width a) (width b) in
      [%e op] (uresize w a) (uresize w b)]
  in*)      

  match expr with
  (* constants ie [123h] or [0b10110h] TODO signed! *)
  | {pexp_desc=Pexp_constant(Pconst_integer(v, Some('h'))); pexp_loc=loc} -> 
    let rec nbits x = 
      if x < 0 then location_exn ~loc "[%hw] constant must be positive";
      match x with 0 | 1 -> 1 | x -> 1 + (nbits (x/2)) 
    in
    let v = int_of_string v in
    [%expr consti [%e int (nbits v)] [%e int v]]

  (* [bit ...] *)
  | [%expr [%e? signal].{[%e? idx]}] ->
    let new_exp = [%expr bit [%e signal] [%e idx]] in
    mapper.expr mapper new_exp 

  (* [select ...] *)
  | [%expr [%e? signal].{[%e? idx0],[%e? idx1]}] ->
    let new_exp = [%expr select [%e signal] [%e idx0] [%e idx1]] in
    mapper.expr mapper new_exp 

  (* map integer to signal operators *)
  | [%expr [%e? a] + [%e? b]] -> mapper.expr mapper [%expr ([%e a] +: [%e b])]
  | [%expr [%e? a] - [%e? b]] -> mapper.expr mapper [%expr ([%e a] -: [%e b])]
  | [%expr [%e? a] * [%e? b]] -> mapper.expr mapper [%expr ([%e a] *: [%e b])]

  | [%expr [%e? a] land [%e? b]] -> mapper.expr mapper [%expr ([%e a] &: [%e b])]
  | [%expr [%e? a] lor  [%e? b]] -> mapper.expr mapper [%expr ([%e a] |: [%e b])]
  | [%expr [%e? a] lxor [%e? b]] -> mapper.expr mapper [%expr ([%e a] ^: [%e b])]
  | [%expr lnot [%e? a]] -> mapper.expr mapper [%expr ~: [%e a]]

  | [%expr [%e? a] <  [%e? b]] -> mapper.expr mapper [%expr ([%e a] <:  [%e b])]
  | [%expr [%e? a] <= [%e? b]] -> mapper.expr mapper [%expr ([%e a] <=: [%e b])]
  | [%expr [%e? a] >  [%e? b]] -> mapper.expr mapper [%expr ([%e a] >:  [%e b])]
  | [%expr [%e? a] >= [%e? b]] -> mapper.expr mapper [%expr ([%e a] >=: [%e b])]

  | [%expr [%e? a] == [%e? b]] -> mapper.expr mapper [%expr ([%e a] ==: [%e b])]
  | [%expr [%e? a] <> [%e? b]] -> mapper.expr mapper [%expr ([%e a] <>: [%e b])]

  (* [if%hw c then t else f] to [mux2 c t f] *)
  | [%expr [%hw if [%e? cond] then [%e? _true] else [%e? _false]]] ->
    let new_exp = [%expr mux2 [%e cond] [%e _true] [%e _false]] in
    mapper.expr mapper new_exp 

  (* [%hw expr] allowable in within [%hw ...] *)
  | [%expr [%hw [%e? e]]] ->
    [%expr [%e hwexpr {mapper with expr=hwexpr} e]]

  | _ -> default_mapper.expr mapper expr

let mapper argv = 
  { default_mapper with

    expr = (fun mapper expr ->
      match expr with
      (* [%hw expr] and 'let%hw pat = <expr> in ...' *)
      | [%expr [%hw [%e? e]]] ->
        [%expr [%e hwexpr {mapper with expr=hwexpr} e]]
      
      | _ -> default_mapper.expr mapper expr);

    structure_item = (fun mapper stri ->
      match stri with
      
      (* [%hw let pat = <expr>] or 'let%hw pat = <expr>' *)
      | [%stri [%%hw let [%p? var] = [%e? e0]]] ->
        [%stri let [%p default_mapper.pat mapper var] = 
          [%e hwexpr {mapper with expr=hwexpr} e0]]

      | _ -> default_mapper.structure_item mapper stri);

  }

let () = register "hardcaml" mapper

Great gist. I massaged the code to look more like yours for comparison: here. The main difference now is that the original code integrate with Ppx_core and Ppx_driver. It is Ppx_core that takes care of recursion over expressions. The ewrap function wrap sub-expressions with [%hw] to ensure that Ppx_core calls the extension mapper on them.