feat: postgresql expression subscripts

ref: https://www.postgresql.org/docs/current/sql-expressions.html#SQL-EXPRESSIONS-SUBSCRIPTS
DerekStride · Dec 11, 2023 · c54eff2 · c54eff2
1 parent 7148a47
commit c54eff2
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 72 deletions.
diff --git a/grammar.js b/grammar.js
@@ -456,78 +456,80 @@ module.exports = grammar({
 
     keyword_array: _ => make_keyword("array"), // not included in _type since it's a constructor literal
 
-    _type: $ => seq(
-      choice(
-        $.keyword_boolean,
-        $.bit,
-        $.binary,
-        $.varbinary,
-        $.keyword_image,
-
-        $.keyword_smallserial,
-        $.keyword_serial,
-        $.keyword_bigserial,
-
-        $.tinyint,
-        $.smallint,
-        $.mediumint,
-        $.int,
-        $.bigint,
-        $.decimal,
-        $.numeric,
-        $.double,
-        $.float,
-
-        $.keyword_money,
-        $.keyword_smallmoney,
-
-        $.char,
-        $.varchar,
-        $.nchar,
-        $.nvarchar,
-        $.numeric,
-        $.keyword_string,
-        $.keyword_text,
-
-        $.keyword_uuid,
-
-        $.keyword_json,
-        $.keyword_jsonb,
-        $.keyword_xml,
-
-        $.keyword_bytea,
-        $.keyword_inet,
-
-        $.enum,
-
-        $.keyword_date,
-        $.keyword_datetime,
-        $.keyword_datetime2,
-        $.datetimeoffset,
-        $.keyword_smalldatetime,
-        $.time,
-        $.timestamp,
-        $.keyword_timestamptz,
-        $.keyword_interval,
-
-        $.keyword_geometry,
-        $.keyword_geography,
-        $.keyword_box2d,
-        $.keyword_box3d,
-
-        $.keyword_oid,
-        $.keyword_name,
-        $.keyword_regclass,
-        $.keyword_regnamespace,
-        $.keyword_regproc,
-        $.keyword_regtype,
-
-        field("custom_type", $.object_reference)
-      ),
-      optional($.array_size_definition)
-    ),
-
-    array_size_definition: $ => seq(
+    _type: $ => prec.left(
+      seq(
+        choice(
+          $.keyword_boolean,
+          $.bit,
+          $.binary,
+          $.varbinary,
+          $.keyword_image,
+
+          $.keyword_smallserial,
+          $.keyword_serial,
+          $.keyword_bigserial,
+
+          $.tinyint,
+          $.smallint,
+          $.mediumint,
+          $.int,
+          $.bigint,
+          $.decimal,
+          $.numeric,
+          $.double,
+          $.float,
+
+          $.keyword_money,
+          $.keyword_smallmoney,
+
+          $.char,
+          $.varchar,
+          $.nchar,
+          $.nvarchar,
+          $.numeric,
+          $.keyword_string,
+          $.keyword_text,
+
+          $.keyword_uuid,
+
+          $.keyword_json,
+          $.keyword_jsonb,
+          $.keyword_xml,
+
+          $.keyword_bytea,
+          $.keyword_inet,
+
+          $.enum,
+
+          $.keyword_date,
+          $.keyword_datetime,
+          $.keyword_datetime2,
+          $.datetimeoffset,
+          $.keyword_smalldatetime,
+          $.time,
+          $.timestamp,
+          $.keyword_timestamptz,
+          $.keyword_interval,
+
+          $.keyword_geometry,
+          $.keyword_geography,
+          $.keyword_box2d,
+          $.keyword_box3d,
+
+          $.keyword_oid,
+          $.keyword_name,
+          $.keyword_regclass,
+          $.keyword_regnamespace,
+          $.keyword_regproc,
+          $.keyword_regtype,
+
+          field("custom_type", $.object_reference)
+        ),
+        optional($.array_size_definition)
+      ),
+    ),
+
+    array_size_definition: $ => prec.left(
       choice(
         seq($.keyword_array, optional($._array_size_definition)),
         repeat1($._array_size_definition),
@@ -3119,6 +3121,7 @@ module.exports = grammar({
         $.exists,
         $.invocation,
         $.binary_expression,
+        $.subscript,
         $.unary_expression,
         $.array,
         $.interval,
@@ -3127,6 +3130,22 @@ module.exports = grammar({
       )
     ),
 
+    subscript: $ => prec.left('binary_is',
+      seq(
+        field('expression', $._expression),
+        "[",
+        choice(
+          field('subscript', $._expression),
+          seq(
+            field('lower', $._expression),
+            ':',
+            field('upper', $._expression),
+          ),
+        ),
+        "]",
+      ),
+    ),
+
     op_other: $ => token(
       choice(
         '->',

diff --git a/test/corpus/expressions.txt b/test/corpus/expressions.txt
@@ -637,3 +637,52 @@ SELECT 'fat cats ate rats' @@ !! ('cat' <-> 'rat'::tsquery);
                       (literal))
                     (object_reference
                       (identifier))))))))))
+
+================================================================================
+Array subscript
+================================================================================
+
+SELECT
+  (ARRAY[1, 4, 9])[1+1],
+  (ARRAY[1, 4, 9])[1:3 - 1],
+  (ARRAY[1, 4, 9])[1:2][3];
+
+--------------------------------------------------------------------------------
+
+(program
+  (statement
+    (select
+      (keyword_select)
+      (select_expression
+        (term
+          value: (subscript
+            expression: (array
+              (keyword_array)
+              (literal)
+              (literal)
+              (literal))
+            subscript: (binary_expression
+              left: (literal)
+              right: (literal))))
+        (term
+          value: (subscript
+            expression: (array
+              (keyword_array)
+              (literal)
+              (literal)
+              (literal))
+            lower: (literal)
+            upper: (binary_expression
+              left: (literal)
+              right: (literal))))
+        (term
+          value: (subscript
+            expression: (subscript
+              expression: (array
+                (keyword_array)
+                (literal)
+                (literal)
+                (literal))
+              lower: (literal)
+              upper: (literal))
+            subscript: (literal)))))))