Source file utils.ml

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
open Core

let simple_query query sexp = Lazy_list.to_list (Semantics.query query sexp)

let get_fields sexp field =
  simple_query (Syntax.Pipe (Syntax.Smash, Syntax.Field field)) sexp
;;

module Non_unique_field = struct
  type t =
    { field : string
    ; sexp : Sexp.t
    ; matches : Sexp.t list
    }
  [@@deriving sexp]
end

let get_one_field sexp field =
  let results = get_fields sexp field in
  match results with
  | [] | _ :: _ :: _ ->
    Or_error.error
      "non-unique field"
      { Non_unique_field.field; sexp; matches = results }
      Non_unique_field.sexp_of_t
  | [ result ] -> Ok result
;;

let sexp_rewrite sexp ~f:visit =
  let rec aux sexp =
    match visit sexp with
    | `Changed sexp' -> sexp'
    | `Unchanged ->
      (match sexp with
       | Sexp.Atom _ -> sexp
       | Sexp.List sexps ->
         let sexps' = List.map ~f:aux sexps in
         if List.for_all2_exn ~f:phys_equal sexps sexps' then sexp else Sexp.List sexps')
  in
  aux sexp
;;

let immediate_fields = function
  | Sexp.List children ->
    List.fold ~init:(Ok []) children ~f:(fun acc child ->
      match acc with
      | Error _ -> acc
      | Ok by_field ->
        (match child with
         | Sexp.List [ Sexp.Atom field; value ] ->
           (match List.Assoc.find by_field ~equal:String.equal field with
            | None -> Ok (List.Assoc.add by_field ~equal:String.equal field value)
            | Some _ -> Or_error.error "multiple values for field" field String.sexp_of_t)
         | _ -> Or_error.error "not a field" child Fn.id))
    |> fun result ->
    (* Restore original order *)
    Or_error.map result ~f:List.rev
  | Sexp.Atom atom -> Or_error.error "not a record" atom String.sexp_of_t
;;

let to_record_sexp by_fields =
  Sexp.List
    (List.map by_fields ~f:(fun (field, value) -> Sexp.List [ Sexp.Atom field; value ]))
;;

let replace_immediate_field ~field ~value sexp =
  Or_error.map (immediate_fields sexp) ~f:(fun by_field ->
    List.Assoc.remove by_field ~equal:String.equal field
    |> (fun by_field -> List.Assoc.add by_field ~equal:String.equal field value)
    |> to_record_sexp)
;;

let replace_field_recursively ~field ~value sexp =
  sexp_rewrite sexp ~f:(function
    | Sexp.List [ Sexp.Atom f; _ ] when String.equal field f ->
      `Changed (Sexp.List [ Sexp.Atom f; value ])
    | _ -> `Unchanged)
;;

let replace_field ~field ~value sexp immediate_or_recursive =
  match immediate_or_recursive with
  | `Immediate -> replace_immediate_field ~field ~value sexp
  | `Recursive ->
    let result = replace_field_recursively ~field ~value sexp in
    if Sexp.( = ) result sexp
    then Or_error.error "field not found" field String.sexp_of_t
    else Ok result
;;

let%test_module "Utils" =
  (module struct
    let sexp =
      Sexp.of_string
        "((first (a b c)) (second 123) (third ()) (fourth ((foo a) (boo b))))"
    ;;

    let%test _ =
      [%compare.equal: Sexp.t Or_error.t]
        (get_one_field sexp "second")
        (Ok (Sexp.Atom "123"))
    ;;

    let%test _ = Result.is_error (get_one_field sexp "zoo")

    let%test _ =
      [%compare.equal: Sexp.t Or_error.t] (get_one_field sexp "boo") (Ok (Sexp.Atom "b"))
    ;;

    let%test _ = Result.is_error (immediate_fields (Sexp.of_string "zoo"))
    let%test _ = Result.is_error (immediate_fields (Sexp.of_string "(zoo)"))
    let%test _ = Result.is_error (immediate_fields (Sexp.of_string "(zoo boo)"))
    let%test _ = Result.is_error (immediate_fields (Sexp.of_string "((good true)(bad))"))

    let%test _ =
      [%equal: Sexp.t]
        (List.Assoc.find_exn
           (Or_error.ok_exn (immediate_fields sexp))
           ~equal:String.equal
           "second")
        (Atom "123")
    ;;

    let%test _ =
      [%equal: Sexp.t]
        (List.Assoc.find_exn
           (Or_error.ok_exn (immediate_fields sexp))
           ~equal:String.equal
           "third")
        (List [])
    ;;

    let%test _ =
      [%equal: Sexp.t]
        (List.Assoc.find_exn
           (Or_error.ok_exn (immediate_fields sexp))
           ~equal:String.equal
           "fourth")
        (Sexp.of_string "((foo a) (boo b))")
    ;;

    let%test _ =
      [%equal: Sexp.t] (to_record_sexp (Or_error.ok_exn (immediate_fields sexp))) sexp
    ;;

    let%test _ =
      let value = Sexp.Atom "my-new-value" in
      let sexp = Or_error.ok_exn (replace_field ~field:"second" ~value sexp `Immediate) in
      [%equal: Sexp.t]
        (List.Assoc.find_exn
           (Or_error.ok_exn (immediate_fields sexp))
           ~equal:String.equal
           "second")
        value
    ;;

    let%test _ =
      let value = Sexp.Atom "my-new-value" in
      let sexp = Or_error.ok_exn (replace_field ~field:"foo" ~value sexp `Recursive) in
      let fourth_value =
        List.Assoc.find_exn
          (Or_error.ok_exn (immediate_fields sexp))
          ~equal:String.equal
          "fourth"
      in
      [%equal: Sexp.t]
        (List.Assoc.find_exn
           (Or_error.ok_exn (immediate_fields fourth_value))
           ~equal:String.equal
           "foo")
        value
    ;;
  end)
;;