Merge pull request #48 from thatportugueseguy/add-array-string-slice

Add array and string slice. update index with tests

Author: davesnx

source/Ast.ml

@@ -50,6 +50,8 @@ type expression =
   | Map of expression
   (* .[] *)
   (* map(x) *)
+  | Slice of int option * int option
+    (* slice(Some 1, Some 10), slice(None, Some 10), slice(Some 1, None) *)
   | FlatMap of expression (* flat_map(x) *)
   | Reduce of expression (* reduce(x) *)
   | Select of expression (* select(x) *)

source/Compiler.ml

@@ -46,9 +46,9 @@ module Output = struct
 
   let lift2 (f : 'a -> 'b -> 'c) (mx : ('a, string) result)
       (my : ('b, string) result) : ('c, string) result =
-    match mx with
-    | Ok x -> ( match my with Ok y -> Ok (f x y) | Error err -> Error err)
-    | Error err -> Error err
+    match (mx, my) with
+    | Ok x, Ok y -> Ok (f x y)
+    | Error err, _ | _, Error err -> Error err
 
   let collect (xs : ('a list, string) result list) : ('a list, string) result =
     List.fold_right (lift2 ( @ )) xs empty
@@ -173,16 +173,58 @@ let index (value : int) (json : Json.t) =
   match json with
   | `List list when List.length list > value ->
       Output.return (Json.index value json)
-  | `List list -> Error (make_acessing_to_missing_item value (List.length list))
+  | `List _ -> Output.return `Null
   | _ -> Error (make_error ("[" ^ string_of_int value ^ "]") json)
 
+let slice (start : int option) (end_ : int option) (json : Json.t) =
+  let start =
+    match (json, start) with
+    | `String s, Some start when start > String.length s -> String.length s
+    | `String s, Some start when start < 0 -> start + String.length s
+    | `List l, Some start when start > List.length l -> List.length l
+    | `List l, Some start when start < 0 -> start + List.length l
+    | (`String _ | `List _), Some start -> start
+    | (`String _ | `List _), None -> 0
+    | _ -> assert false
+  in
+  let end_ =
+    match (json, end_) with
+    | `String s, None -> String.length s
+    | `String s, Some end_ when end_ > String.length s -> String.length s
+    | `String s, Some end_ when end_ < 0 -> end_ + String.length s
+    | `List l, None -> List.length l
+    | `List l, Some end_ when end_ > List.length l -> List.length l
+    | `List l, Some end_ when end_ < 0 -> end_ + List.length l
+    | (`String _ | `List _), Some end_ -> end_
+    | _ -> assert false
+  in
+  match json with
+  | `String _s when end_ < start -> Output.return (`String "")
+  | `String s -> Output.return (`String (String.sub s start (end_ - start)))
+  | `List _l when end_ < start -> Output.return (`List [])
+  | `List l ->
+      let sliced =
+        List.fold_left
+          (fun (acc, i) x ->
+            if i >= start && i < end_ then (x :: acc, i + 1) else (acc, i + 1))
+          ([], 0) l
+        |> fst |> List.rev
+      in
+      Output.return (`List sliced)
+  | _ ->
+      Error
+        (make_error
+           ("[" ^ string_of_int start ^ ":" ^ string_of_int end_ ^ "]")
+           json)
+
 let rec compile expression json : (Json.t list, string) result =
   match expression with
   | Identity -> Output.return json
   | Empty -> empty
   | Keys -> keys json
   | Key (key, opt) -> member key opt json
   | Index idx -> index idx json
+  | Slice (start, end_) -> slice start end_ json
   | Head -> head json
   | Tail -> tail json
   | Length -> length json

source/Parser.mly

@@ -28,6 +28,7 @@
 %token CLOSE_BRACKET
 
 %token COMMA
+%token COLON
 /* %token OPEN_BRACE
 %token CLOSE_BRACE */
 %token EOF
@@ -92,20 +93,26 @@ item_expr:
   | e = term
     { e }
 
+number:
+  | n = NUMBER;
+    { n }
+  | SUB; n = NUMBER;
+    { -.n }
+
 term:
   | DOT;
     { Identity }
   | RECURSE;
     { Recurse }
   | s = STRING;
     { Literal (String s) }
-  | n = NUMBER;
+  | n = number;
     { Literal (Number n) }
   | b = BOOL;
     { Literal (Bool b) }
   | NULL
     { Literal(Null) }
-  | f = FUNCTION; from = NUMBER; SEMICOLON; upto = NUMBER; CLOSE_PARENT;
+  | f = FUNCTION; from = number; SEMICOLON; upto = number; CLOSE_PARENT;
     { match f with
       | "range" -> Range (int_of_float from, int_of_float upto)
       | _ -> failwith (f ^ " is not a valid function")
@@ -197,9 +204,21 @@ term:
   // Parentheses allow a full sequence_expr inside, reducing to an item_expr
   | OPEN_PARENT; e = sequence_expr; CLOSE_PARENT;
     { e }
-  | e = term; OPEN_BRACKET; i = NUMBER; CLOSE_BRACKET
+  | e = term; OPEN_BRACKET; i = number; CLOSE_BRACKET
     { Pipe (e, Index (int_of_float i)) }
 
+  /* Full slice with both indices: .[1:5] */
+  | e = term; OPEN_BRACKET; start = number; COLON; end_ = number; CLOSE_BRACKET
+    { Pipe (e, Slice (Some (int_of_float start), Some (int_of_float end_))) }
+
+  /* Start-only slice: .[3:] */
+  | e = term; OPEN_BRACKET; start = number; COLON; CLOSE_BRACKET
+    { Pipe (e, Slice (Some (int_of_float start), None)) }
+
+  /* End-only slice: .[:3] */
+  | e = term; OPEN_BRACKET; COLON; end_ = number; CLOSE_BRACKET
+    { Pipe (e, Slice (None, Some (int_of_float end_))) }
+
   | DOT; k = STRING; opt = boption(QUESTION_MARK)
     { Key (k, opt) }
 

source/Tokenizer.ml

@@ -18,6 +18,7 @@ type token =
   | OPEN_BRACKET
   | CLOSE_BRACKET
   | SEMICOLON
+  | COLON
   | DOT
   | RECURSE
   | PIPE
@@ -80,6 +81,7 @@ let rec tokenize buf =
   | "]" -> Ok CLOSE_BRACKET
   | "|" -> Ok PIPE
   | ";" -> Ok SEMICOLON
+  | ":" -> Ok COLON
   | "," -> Ok COMMA
   | "?" -> Ok QUESTION_MARK
   | "null" -> Ok NULL

test/Test_parse.ml

@@ -19,6 +19,7 @@ let case input expected =
 
 let tests =
   [
+    case ".[-1]" (Pipe (Identity, Index (-1)));
     case ".[1]" (Pipe (Identity, Index 1));
     case "[1]" (List [ Literal (Number 1.) ]);
     case ".store.books" (Pipe (Key ("store", false), Key ("books", false)));
@@ -44,4 +45,8 @@ let tests =
     case "[1][0]" (Pipe (List [ Literal (Number 1.) ], Index 0));
     case "[1].foo" (Pipe (List [ Literal (Number 1.) ], Key ("foo", false)));
     case "(empty).foo?" (Pipe (Empty, Key ("foo", true)));
+    case ".[1:3]" (Pipe (Identity, Slice (Some 1, Some 3)));
+    case ".[1:]" (Pipe (Identity, Slice (Some 1, None)));
+    case ".[:3]" (Pipe (Identity, Slice (None, Some 3)));
+    case ".[-2:]" (Pipe (Identity, Slice (Some (-2), None)));
   ]

test/Test_runtime.ml

@@ -30,6 +30,30 @@ let tests =
     test "[[2]]" "[3]" "[ [ 2 ] ]";
     (* test "[{}]" "[2]" "[{}]"; *)
     (* test "[.[]]" "[\"a\"]" "[\"a\"]"; *)
+
+    (* Array index tests *)
+    test ".[0]" {|["a","b","c","d","e"]|} {|"a"|};
+    test ".[3]" {|["a","b"]|} "null";
+    test ".[-1]" {|["a","b","c","d","e"]|} {|"e"|};
+
+    (* Array / String slice tests *)
+    test ".[2:4]" {|["a","b","c","d","e"]|} {|[ "c", "d" ]|};
+    test ".[2:4]" {|"abcdefghi"|} {|"cd"|};
+    test ".[5:7]" {|["a","b","c"]|} {|[]|};
+    test ".[5:7]" {|"abc"|} {|""|};
+    test ".[:3]" {|["a","b","c","d","e"]|} {|[ "a", "b", "c" ]|};
+    test ".[:-2]" {|["a","b","c","d","e"]|} {|[ "a", "b", "c" ]|};
+    test ".[:3]" {|"abcdefghi"|} {|"abc"|};
+    test ".[:-2]" {|"abcdefghi"|} {|"abcdefg"|};
+    test ".[-2:]" {|["a","b","c","d","e"]|} {|[ "d", "e" ]|};
+    test ".[2:]" {|["a","b","c","d","e"]|} {|[ "c", "d", "e" ]|};
+    test ".[-2:]" {|"abcdefghi"|} {|"hi"|};
+    test ".[2:]" {|"abcdefghi"|} {|"cdefghi"|};
+    test ".[-4:-2]" {|["a","b","c","d","e"]|} {|[ "b", "c" ]|};
+    test ".[-2:-4]" {|["a","b","c","d","e"]|} {|[]|};
+    test ".[-4:-2]" {|"abcdefghi"|} {|"fg"|};
+    test ".[-2:-4]" {|"abcde"|} {|""|};
+
     test "1,1" "[]" "1\n1";
     test "1,." "[]" "1\n[]";
     test {|.foo | .bar|} {|{"foo": {"bar": 42}, "bar": "badvalue"}|} {|42|};