The monadic interface of snarky

[mimoo]

Snarky exposes a functor that can be used to instantiate its API in src/snark0.ml. In the file, there's actually two functors: one for the monadic API (Make), and one for the imperative API (Run.Make).

See here for a short explainer on monads.

You can see the monadic API being instantiated like so:

module Impl = Snarky_backendless.Snark.Make (struct
    include Kimchi_backend.Pasta.Vesta_based_plonk
    module Inner_curve = Kimchi_backend.Pasta.Pasta.Pallas
  end

and the imperative API being instantiated like so:

module Impl = Snarky_backendless.Snark.Run.Make (struct
  include Kimchi_backend.Pasta.Vesta_based_plonk
  module Inner_curve = Kimchi_backend.Pasta.Pasta.Pallas
end)

Note: it can get confused as the imperative and monadic APIs are not clearly distinct in the code. The word "Run" seems to be used in some places to mean "imperative API", and in some places to mean "compilation of the AST".

The imperative API is a thin wrapper built on top of the monadic one. In this document I explain how the monadic API works (which really is most of snarky).

The AST of the Monadic API

The monadic API relies on an AST defined in checked_ast.ml:

type ('a, 'f) t =
    | Pure : 'a -> ('a, 'f) t
    | Direct :
        ('f Run_state.t -> 'f Run_state.t * 'a) * ('a -> ('b, 'f) t)
        -> ('b, 'f) t
    | Add_constraint : ('f Cvar.t, 'f) Constraint.t * ('a, 'f) t -> ('a, 'f) t
    | As_prover : (unit, 'f) Types.As_prover.t * ('a, 'f) t -> ('a, 'f) t
    | Lazy : ('a, 'f) t * ('a Lazy.t -> ('b, 'f) t) -> ('b, 'f) t
    | With_label : string * ('a, 'f) t * ('a -> ('b, 'f) t) -> ('b, 'f) t
    | With_handler :
        Request.Handler.single * ('a, 'f) t * ('a -> ('b, 'f) t)
        -> ('b, 'f) t
    | Exists :
        ('var, 'value, 'f, (unit, 'f) t) Types.Typ.t
        * ( ('value Request.t, 'f) Types.As_prover.t
          , ('value, 'f) Types.As_prover.t )
          Provider.t
        * (('var, 'value) Handle.t -> ('a, 'f) t)
        -> ('a, 'f) t
    | Next_auxiliary : (int -> ('a, 'f) t) -> ('a, 'f) t

Each variant is an AST node, where ('a, 'f) t means that it returns a value of type 'a (and 'f is the field used). For example:

| Add_constraint : ('f Cvar.t, 'f) Constraint.t * ('a, 'f) t -> ('a, 'f) t

represents a constraint being added. It contains a tuple of two children: the constraint we want to add, and the rest of the AST (which can be used to compute a value of type 'a).

Note: the AST is not constructed like the ASTs I'm typically used to, which build a tree bottom up. Instead, the AST looks more like a linked list, and each operation appends to that list.

To read this AST, one important thing to know is that As_prover.t is just the type signature of a closure that takes a map from variables to their values, and return a value:

('f Cvar.t -> 'f) -> 'a

So As_prover.t closures are ran only when we have access to the actual values of variables: either during debugging, or during witness generation.

TODO: it would be nice to have comments on each node of the AST to understand exactly their purpose. Some of the nodes are lost on me.

There's two aspects to using the monadic API:

1. construct a circuit by building the AST
2. running through the AST to compile the circuit

The rest of this document goes over both of these.

Constructing a circuit with the monadic API

You can see how a circuit is constructed in the tests:

module Impl = Snarky_backendless.Snark.Make (struct
  include Kimchi_backend.Pasta.Vesta_based_plonk
  module Inner_curve = Kimchi_backend.Pasta.Pasta.Pallas
end)

let main ((b1, b2) : Impl.Boolean.var * Impl.Boolean.var) =
  let open Impl in
  let%bind x = exists Boolean.typ ~compute:(As_prover.return true) in
  let%bind y = exists Boolean.typ ~compute:(As_prover.return true) in
  let%bind z = Boolean.(x &&& y) in
  let%bind b3 = Boolean.(b1 && b2) in

  let%bind () = Boolean.Assert.is_true z in
  let%bind () = Boolean.Assert.is_true b3 in

  Checked.return ()

let circuit_hash =
  let input_typ = Impl.Typ.tuple2 Impl.Boolean.typ Impl.Boolean.typ in
  let return_typ = Impl.Typ.unit in
  let cs : Impl.R1CS_constraint_system.t =
    Impl.constraint_system ~input_typ ~return_typ main
  in
  Md5.to_hex (Impl.R1CS_constraint_system.digest cs)

Notice that ppx_let is use with the let% syntax.

These functions are all defined in checked_ast.ml in a Basic module:

module Basic :
  Checked_intf.Basic with module Types = Types with type 'f field = 'f = struct
  module Types = Types

  type ('a, 'f) t = ('a, 'f) Types.Checked.t

  type 'f field = 'f

  include Monad_let.Make2 (T0)

  (* see, for example, add_constraint here *)
  let add_constraint c = Add_constraint (c, return ())

  let as_prover x = As_prover (x, return ())

The functions are quite simple there and mostly construct nodes of the tree (or extend the tree via bind).
For example, you can see that exists will simply create a new node, with an empty return instruction which can be used later to extend the program.

let exists typ p = Exists (typ, p, return)

The logic wrapper by a functor in checked.ml, which adds some functionality around computing witness values (with some provider/handler logic, which we will not explain here):

module Make
    (Basic : Checked_intf.Basic)
    (As_prover : As_prover_intf.Basic
                   with type ('a, 'f) t = ('a, 'f) Basic.Types.As_prover.t
                    and type 'f field := 'f Basic.field) :
  Checked_intf.S
    with module Types = Basic.Types
    with type 'f field = 'f Basic.field = struct
  include Basic

  let request_witness (typ : ('var, 'value, 'f field) Types.Typ.t)
      (r : ('value Request.t, 'f field) As_prover.t) =
    let%map h = exists typ (Request r) in
    Handle.var h

  let request ?such_that typ r =
    match such_that with
    | None ->
        request_witness typ (As_prover.return r)
    | Some such_that ->
        let open Let_syntax in
        let%bind x = request_witness typ (As_prover.return r) in
        let%map () = such_that x in
        x

  (* see, for example, how `exists` is wrapped *)
  let exists_handle ?request ?compute typ =
    let provider =
      let request =
        Option.value request ~default:(As_prover.return Request.Fail)
      in
      match compute with
      | None ->
          Types.Provider.Request request
      | Some c ->
          Types.Provider.Both (request, c)
    in
    exists typ provider

  let exists ?request ?compute typ =
    let%map h = exists_handle ?request ?compute typ in
    Handle.var h

This is it for the construction of the AST.
Next let's look at how snarky parses the AST.

Compiling a circuit with the monadic API

Compiling a circuit with the monadic API is call "running" it. Lots of functions are called run so beware.

I suspect this is because a monad is used, and the function to get out of a monad is often called run.

The runner is created in checked_runner.ml by a Make functor.
This Make functor essentially calls another functor in ast_runner.ml (Make).

The ast_runner.ml functor looks like this:

module Make_runner
    (Checked : Checked_intf.Basic
                 with type ('a, 'f) Types.As_prover.t =
                   ('a, 'f) Types.As_prover.t
                  and type ('a, 'f) Types.Provider.provider =
                   ('a, 'f) Types.Provider.provider) =
struct
  let handle_error label f =
    try f () with
    | Runtime_error (stack, exn, bt) ->
        let bt_new = Printexc.get_backtrace () in
        raise (Runtime_error (label :: stack, exn, bt ^ "\n\n" ^ bt_new))
    | exn ->
        let bt = Printexc.get_backtrace () in
        raise (Runtime_error ([ label ], exn, bt))

  let rec run :
      type a.
      (a, 'f Checked.field) Checked_ast.t -> (a, 'f Checked.field) Checked.t =
   fun t ->
    match t with
    | As_prover (x, k) ->
        Checked.bind (Checked.as_prover x) ~f:(fun () -> run k)
    | Pure x ->
        Checked.return x
    | Direct (d, k) ->
        Checked.bind (Checked.direct d) ~f:(fun y -> run (k y))
    | Lazy (x, k) ->
        Checked.bind (Checked.mk_lazy (fun () -> run x)) ~f:(fun y -> run (k y))
    | With_label (lab, t, k) ->
        Checked.bind
          (Checked.with_label lab (fun () -> handle_error lab (fun () -> run t)))
          ~f:(fun y -> run (k y))
    | Add_constraint (c, t) ->
        Checked.bind (Checked.add_constraint c) ~f:(fun () -> run t)
    | With_handler (h, t, k) ->
        Checked.bind
          (Checked.with_handler h (fun () -> run t))
          ~f:(fun y -> run (k y))
    | Exists
        ( Typ
            { var_to_fields
            ; var_of_fields
            ; value_to_fields
            ; value_of_fields
            ; size_in_field_elements
            ; constraint_system_auxiliary
            ; check
            }
        , p
        , k ) ->
        let typ =
          Types.Typ.Typ
            { var_to_fields
            ; var_of_fields
            ; value_to_fields
            ; value_of_fields
            ; size_in_field_elements
            ; constraint_system_auxiliary
            ; check = (fun var -> run (check var))
            }
        in
        Checked.bind (Checked.exists typ p) ~f:(fun y -> run (k y))
    | Next_auxiliary k ->
        Checked.bind (Checked.next_auxiliary ()) ~f:(fun y -> run (k y))
end

Interestingly, we are using a state monad to parse the AST.

Note: my guess is that it's a natural way to parse a linked-list-like AST. I wrote more about state monads here.

The actual logic behind the Checked function that can handle the different nodes in the tree is implemented in checked_runner.ml, and the two are glued together in a functor inside of checked_runner.ml as well.

The flow is a bit convoluted, but it looks like this:

• Ast_runner.Make_runner relies on a Checked_intf.Basic type signature
• inside Checked_runner.Make, the basic checked interface is created by the Checked_runner.Make_checked functor
• still inside Checked_runner.Make, the newly created basic checked interface is passed to Ast_runner.Make_runer

to summarize this in pseudo-code:

Checked_runner.Make():
	module Checked = Checked_runner.Make_checked()
	include Ast_runner.Make_runner(Checked)

More importantly, the state monad has the following type signature:

type ('a, 'f) t = run_state -> run_state * 'a

Which performs a state transition on a Run_state.t and produces a value 'a'.

So the function that runs the AST actually parses the AST with the state monad, to create such a function. At the end, we are left with a big state transition acting on a (presumably initial) Run_state.t, which is chaining multiple state transitions together. This represents the computation of the AST.

What is left is just to run that function with an initialized state (and parameterized the state depending on if we want to compile the circuit, generate a witness, or do both).

[bkase]

Thanks @mimoo ! I'll share a few of my thoughts on DSLs because it's relevant I think!

Note: the AST is not constructed like the ASTs I'm typically used to, which build a tree bottom up. Instead, the AST looks more like a linked list, and each operation appends to that list.

A few of the nodes are tree-like -- Lazy, With_label, and Exists
I say "tree-like" because part of it is under a function binding, so it's not strictly "listy" I'd say

let exists typ p = Exists (typ, p, return)

This pattern is a typical way to work with "initial-style monadic embedded DSLs" in functional languages.

This is in contrast to a "final DSL" which uses interfaces to switch on the implementation rather than reinterpreting AST structures.
In Rust, a final DSL would be one where you define all the steps as functions on a trait and then you implement the trait in different ways to get different interpreters.

An initial DSL you'd create by making a recursive Rust enum and to interpret you write a function that recursively traverses the AST.

I think of it like this: an initial DSL reifies the AST into a data structure -- in its abstract form -- that you can touch and feel and actually represent in memory as bytes.

While a final DSL's AST never really exists in memory, it's a shadow of its execution trace -- you can still interpret a final DSL into some serialized form, but that's just implemented as a concrete interpretation of the abstract structure.