1 files changed, 176 insertions, 0 deletions

diff --git a/lib/nix/printer.ml b/lib/nix/printer.ml
new file mode 100644
index 0000000..57e81f4
--- /dev/null
+++ b/lib/nix/printer.ml
@@ -0,0 +1,176 @@
+open Core
+open Types
+open PPrint
+
+let rec escape_chlist = function
+  | [] -> []
+  | '$' :: '{' :: l' -> '\\' :: '$' :: '{' :: escape_chlist l'
+  | '\n' :: l' -> '\\' :: 'n' :: escape_chlist l'
+  | '\r' :: l' -> '\\' :: 'r' :: escape_chlist l'
+  | '\t' :: l' -> '\\' :: 't' :: escape_chlist l'
+  | '\\' :: l' -> '\\' :: '\\' :: escape_chlist l'
+  | '"' :: l' -> '\\' :: '"' :: escape_chlist l'
+  | c :: l' -> c :: escape_chlist l'
+
+let escape_string s = s |> String.to_list |> escape_chlist |> String.of_list
+let out_width = ref 80
+let set_width i = out_width := i
+let indent = ref 2
+let set_indent i = indent := i
+
+let rec doc_of_expr = function
+  | BinaryOp (op, lhs, rhs) ->
+      let lhs_doc = maybe_parens_bop op `Left lhs
+      and rhs_doc = maybe_parens_bop op `Right rhs in
+      infix !indent 1 (doc_of_bop op) lhs_doc rhs_doc
+  | UnaryOp (op, e) -> precede (doc_of_uop op) (maybe_parens (prec_of_uop op) e)
+  | Cond (e1, e2, e3) ->
+      surround !indent 1
+        (soft_surround !indent 1 (string "if") (doc_of_expr e1) (string "then"))
+        (doc_of_expr e2)
+        (string "else" ^^ nest !indent (break 1 ^^ doc_of_expr e3))
+  | With (e1, e2) ->
+      flow (break 1) [ string "with"; doc_of_expr e1 ^^ semi; doc_of_expr e2 ]
+  | Assert (e1, e2) ->
+      flow (break 1) [ string "assert"; doc_of_expr e1 ^^ semi; doc_of_expr e2 ]
+  | Test (e, path) ->
+      maybe_parens 4 e ^^ space ^^ string "?"
+      ^^ group (break 1 ^^ separate_map dot doc_of_expr path)
+  | SetLet bs ->
+      surround !indent 1
+        (string "let " ^^ lbrace)
+        (group (separate_map (break 1) doc_of_binding bs))
+        rbrace
+  | Let (bs, e) ->
+      surround !indent 1 (string "let")
+        (separate_map (break 1) doc_of_binding bs)
+        (prefix !indent 1 (string "in") (doc_of_expr e))
+  | Val v -> doc_of_val v
+  | Id id -> string id
+  | Select (e, path, oe) ->
+      maybe_parens 1 e ^^ dot ^^ doc_of_attpath path
+      ^^ optional
+           (fun e ->
+             space ^^ string "or" ^^ nest !indent (break 1 ^^ maybe_parens 1 e))
+           oe
+  | Apply (e1, e2) -> prefix !indent 1 (maybe_parens 2 e1) (maybe_parens 2 e2)
+  | Aquote e -> surround !indent 0 (string "${") (doc_of_expr e) (string "}")
+
+and maybe_parens lvl e =
+  if prec_of_expr e >= lvl then surround !indent 0 lparen (doc_of_expr e) rparen
+  else doc_of_expr e
+
+and maybe_parens_bop op (loc : [ `Left | `Right ]) e =
+  match (loc, assoc_of_bop op) with
+  | (`Left, Some Left | `Right, Some Right)
+    when prec_of_expr e >= prec_of_bop op ->
+      doc_of_expr e
+  | _, _ -> maybe_parens (prec_of_bop op) e
+
+and doc_of_attpath path = separate_map dot doc_of_expr path
+
+and doc_of_paramset (params, kind) =
+  let ps =
+    List.map ~f:doc_of_param params
+    @ if Poly.(kind = Loose) then [ string "..." ] else []
+  in
+  surround !indent 0 lbrace (separate (comma ^^ break 1) ps) rbrace
+
+and doc_of_param (id, oe) =
+  string id ^^ optional (fun e -> qmark ^^ space ^^ doc_of_expr e) oe
+
+and doc_of_binding = function
+  | AttrPath (path, e) ->
+      doc_of_attpath path ^^ space ^^ equals ^^ space ^^ doc_of_expr e ^^ semi
+  | Inherit (oe, ids) ->
+      let id_docs =
+        List.map
+          ~f:(function
+            | Id x | Val (Str (x, [])) -> string x | _ -> assert false)
+          ids
+      in
+      let xs =
+        flow (break 1)
+          (match oe with
+          | Some e -> parens (doc_of_expr e) :: id_docs
+          | None -> id_docs)
+      in
+      soft_surround !indent 0 (string "inherit" ^^ space) xs semi
+
+and doc_of_bop = function
+  | Plus -> plus
+  | Minus -> minus
+  | Mult -> star
+  | Div -> slash
+  | Gt -> rangle
+  | Lt -> langle
+  | Lte -> string "<="
+  | Gte -> string ">="
+  | Eq -> string "=="
+  | Neq -> string "!="
+  | Or -> string "||"
+  | And -> string "&&"
+  | Impl -> string "->"
+  | Merge -> string "//"
+  | Concat -> string "++"
+
+and doc_of_uop = function Negate -> minus | Not -> bang
+
+and doc_of_val = function
+  | Str (start, xs) ->
+      dquotes
+        (string (escape_string start)
+        ^^ concat
+             (List.map
+                ~f:(fun (e, s) ->
+                  surround !indent 0 (string "${") (doc_of_expr e)
+                    (string "}" ^^ string (escape_string s)))
+                xs))
+  | IStr (i, start, xs) ->
+      let qq = string "''" in
+      let str s =
+        String.split ~on:'\n' s
+        |> List.map ~f:(fun s ->
+               let len = String.length s in
+               let s' =
+                 if len >= i then String.sub s ~pos:i ~len:(len - i) else s
+               in
+               string s')
+        |> separate hardline
+      in
+      enclose qq qq
+        (str start
+        ^^ concat
+             (List.map
+                ~f:(fun (e, s) ->
+                  enclose (string "${") rbrace (doc_of_expr e) ^^ str s)
+                xs))
+  | Int x | Float x | Path x | SPath x | HPath x | Uri x -> string x
+  | Lambda (pattern, body) ->
+      let pat =
+        match pattern with
+        | Alias id -> string id
+        | ParamSet (None, ps) -> doc_of_paramset ps
+        | ParamSet (Some id, ps) ->
+            doc_of_paramset ps ^^ group (break 1 ^^ at ^^ break 1 ^^ string id)
+      in
+      flow (break 1) [ pat ^^ colon; doc_of_expr body ]
+  | List [] -> lbracket ^^ space ^^ rbracket
+  | List es ->
+      surround !indent 1 lbracket
+        (separate_map (break 1) (maybe_parens 2) es)
+        rbracket
+  | AttSet (Nonrec, []) -> lbrace ^^ space ^^ rbrace
+  | AttSet (Nonrec, bs) ->
+      surround !indent 1 lbrace
+        (group (separate_map (break 1) doc_of_binding bs))
+        rbrace
+  | AttSet (Rec, bs) ->
+      string "rec" ^^ space ^^ doc_of_val (AttSet (Nonrec, bs))
+
+let print chan expr = ToChannel.pretty 0.7 !out_width chan (doc_of_expr expr)
+
+let to_string expr =
+  let buf = Stdlib.Buffer.create 0 in
+  ToBuffer.pretty 0.7 !out_width buf (doc_of_expr expr);
+  Stdlib.Buffer.contents buf