aboutsummaryrefslogtreecommitdiff
path: root/rego-proofs/compiler-rewrite/lib/ast.ml
blob: 052496d1e522ce6ed5c32afbbe2b68874995b393 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
open Core
open Yojson.Basic.Util

(* AST types. *)

type expr = Term of term | Every of string option * string * term * query
[@@deriving eq, show]

and literal =
  | Exists of expr list * with_modifier option
  | Expr of expr * with_modifier option
  | Negate of expr * with_modifier option
[@@deriving eq, show]

and query = literal list [@@deriving eq, show]

and term =
  | Scalar of scalar
  | Var of string
  | Ref of reference
  | Array of term list
  | Object of (object_key * term) list
  | Set of term list
  | ArrayComprehension of term * query
  | ObjectComprehension of (object_key * term) * query
  | SetComprehension of term * query
  | Call of reference * term list
[@@deriving eq, show]

and scalar = String of string | Number of float | Bool of bool | Null
[@@deriving eq, show]

and reference = ref_target * term list [@@deriving eq, show]

and ref_target =
  | Var of string
  | Array of term list
  | Object of (object_key * term) list
  | Set of term list
  | ArrayComprehension of term * query
  | ObjectComprehension of (object_key * term) * query
  | SetComprehension of term * query
  | Call of reference * term list
[@@deriving eq, show]

and object_key = Scalar of scalar | Ref of reference | Var of string
[@@deriving eq, show]

and with_modifier = (term * term) list [@@deriving eq, show]

type assign_operator = Assign | Unify [@@deriving eq, show]
type rule_head_lhs = reference [@@deriving eq, show]

type rule_head_rhs =
  | Assignment of assign_operator * term
  | Function of assign_operator * string list * term
  | Set of term
[@@deriving eq, show]

type rule_head = rule_head_lhs * rule_head_rhs [@@deriving eq, show]

type rule = {
  default : bool;
  head : rule_head;
  body : query;
  otherwise : rule option;
}
[@@deriving eq, show]

type import = { alias : string option; path : reference } [@@deriving eq, show]

type package = { path : reference; imports : import list; policy : rule list }
[@@deriving eq, show]

(* Additional JSON utility functions. *)

let is_null : Yojson.Basic.t -> bool = function `Null -> true | _ -> false

let to_bool_or_false : Yojson.Basic.t -> bool = function
  | `Bool b -> b
  | _ -> false

let to_list_or_empty : Yojson.Basic.t -> Yojson.Basic.t list = function
  | `Null -> []
  | `List xs -> xs
  | _ -> failwith "Not a list or null"

let object_key_of_term : term -> object_key = function
  | Scalar s -> Scalar s
  | Ref r -> Ref r
  | Var v -> Var v
  | _ -> failwith "failed to convert term to object_key"

(* Code to parse the JSON output of `opa parse`. *)

let rec parse_term (tree : Yojson.Basic.t) : term =
  let term_type = tree |> member "type" |> to_string in
  match term_type with
  | "var" -> Var (tree |> member "value" |> to_string)
  | "ref" -> Ref (tree |> member "value" |> to_list |> parse_ref)
  | "boolean" -> Scalar (Bool (tree |> member "value" |> to_bool))
  | "number" -> Scalar (Number (tree |> member "value" |> to_number))
  | "string" -> Scalar (String (tree |> member "value" |> to_string))
  | "array" ->
      Array (tree |> member "value" |> to_list |> List.map ~f:parse_term)
  | "set" -> Set (tree |> member "value" |> to_list |> List.map ~f:parse_term)
  | "object" ->
      Object
        (tree |> member "value" |> to_list
        |> List.map ~f:(function
             | `List [ x; y ] ->
                 (x |> parse_term |> object_key_of_term, y |> parse_term)
             | _ -> failwith "ill-formed object"))
  | "arraycomprehension" ->
      ArrayComprehension
        ( tree |> member "value" |> member "term" |> parse_term,
          tree |> member "value" |> member "body" |> to_list
          |> List.map ~f:parse_literal )
  | "setcomprehension" ->
      SetComprehension
        ( tree |> member "value" |> member "term" |> parse_term,
          tree |> member "value" |> member "body" |> to_list
          |> List.map ~f:parse_literal )
  | "objectcomprehension" ->
      ObjectComprehension
        ( ( tree |> member "value" |> member "key" |> parse_term
            |> object_key_of_term,
            tree |> member "value" |> member "value" |> parse_term ),
          tree |> member "value" |> member "body" |> to_list
          |> List.map ~f:parse_literal )
  | "call" -> (
      match tree |> member "value" |> to_list with
      | func :: args ->
          Call
            ( func |> member "value" |> to_list |> parse_ref,
              List.map ~f:parse_term args )
      | _ -> failwith "ill-formed call")
  | _ -> failwith ("unsupported term type " ^ term_type)

and parse_literal tree : literal =
  let with_expr = tree |> member "with" in
  let with_expr =
    match with_expr with
    | `Null -> None
    | `List xs ->
        Some
          (List.map xs ~f:(fun x ->
               ( x |> member "target" |> parse_term,
                 x |> member "value" |> parse_term )))
    | _ ->
        failwith
          ("ill-formed with expression " ^ Yojson.Basic.to_string with_expr)
  in
  match tree |> member "terms" with
  | `List (func :: args) ->
      Expr
        ( Term
            (Call
               ( func |> member "value" |> to_list |> parse_ref,
                 List.map ~f:parse_term args )),
          with_expr )
  | `List _ -> failwith "ill-formed call"
  | x ->
      if not (x |> member "symbols" |> is_null) then
        Exists
          ( x |> member "symbols" |> to_list
            |> List.map ~f:(fun x -> Term (parse_term x)),
            with_expr )
      else if x |> member "domain" |> is_null then
        Expr (Term (parse_term x), with_expr)
      else
        Expr
          ( Every
              ( (if x |> member "key" |> is_null then None
                 else Some (x |> member "key" |> member "value" |> to_string)),
                x |> member "value" |> member "value" |> to_string,
                x |> member "domain" |> parse_term,
                List.map ~f:parse_literal (x |> member "body" |> to_list) ),
            with_expr )

and parse_ref_target tree : ref_target =
  let term = parse_term tree in
  match term with
  | Var term -> Var term
  | Array term -> Array term
  | Object term -> Object term
  | Set term -> Set term
  | ArrayComprehension (term, y) -> ArrayComprehension (term, y)
  | ObjectComprehension (term, y) -> ObjectComprehension (term, y)
  | SetComprehension (term, y) -> SetComprehension (term, y)
  | Call (term, y) -> Call (term, y)
  | _ -> failwith ("unsupported ref_target " ^ show_term term)

and parse_ref tree =
  match tree with
  | target :: args -> (parse_ref_target target, List.map ~f:parse_term args)
  | _ -> failwith "empty reference"

let parse_head tree =
  (* Variables are desugared to empty references. *)
  let lhs = tree |> member "ref" |> to_list |> parse_ref in
  let assign_operator =
    if tree |> member "assign" |> to_bool_or_false then Assign else Unify
  in
  let rhs =
    if not (tree |> member "key" |> is_null) then
      Set (tree |> member "key" |> parse_term)
    else if tree |> member "args" |> to_bool_or_false then
      Function
        ( assign_operator,
          tree |> member "args" |> to_list
          |> List.map ~f:(fun x -> x |> member "value" |> to_string),
          tree |> member "value" |> parse_term )
    else Assignment (assign_operator, tree |> member "value" |> parse_term)
  in
  (lhs, rhs)

let rec parse_rule tree =
  {
    default = tree |> member "default" |> to_bool_or_false;
    head = tree |> member "head" |> parse_head;
    body = List.map ~f:parse_literal (tree |> member "body" |> to_list_or_empty);
    otherwise =
      (match tree |> member "else" with
      | `Null -> None
      | else_branch -> Some (parse_rule else_branch));
  }

let parse_import_path tree =
  match tree |> member "type" |> to_string with
  | "ref" -> tree |> member "value" |> to_list |> parse_ref
  | _ -> failwith "unimplemented"

let parse_import tree =
  {
    alias = tree |> member "alias" |> to_string_option;
    path = tree |> member "path" |> parse_import_path;
  }

let parse_package tree = tree |> member "path" |> to_list |> parse_ref

let parse_tree tree =
  {
    path = tree |> member "package" |> parse_package;
    imports =
      List.map ~f:parse_import (tree |> member "imports" |> to_list_or_empty);
    policy = List.map ~f:parse_rule (tree |> member "rules" |> to_list_or_empty);
  }

(* OPA interaction. *)

let invoke ?(stdin = "") ?(env = [||]) cmd =
  let ic, oc, ec = UnixLabels.open_process_full cmd ~env in
  Out_channel.output_string oc stdin;
  Out_channel.close oc;
  (In_channel.input_all ic, In_channel.input_all ec)

let ast_of_string policy_code =
  let stdout, _ =
    invoke "/home/jakob/.local/bin/opa parse -f json /dev/stdin"
      ~stdin:policy_code ~env:[||]
  in
  Yojson.Basic.from_string stdout |> parse_tree

© 2015 - 2026 Jakob L. Kreuze