Add MultiRange

edgedb.toml

@@ -1,2 +1,2 @@
 [edgedb]
-server-version = "3.0"
+server-version = "4.0"

src/driver/src/main/java/com/edgedb/driver/binary/codecs/MultiRangeCodec.java

@@ -0,0 +1,56 @@
+package com.edgedb.driver.binary.codecs;
+
+import com.edgedb.driver.binary.PacketReader;
+import com.edgedb.driver.binary.PacketWriter;
+import com.edgedb.driver.binary.protocol.common.descriptors.CodecMetadata;
+import com.edgedb.driver.datatypes.MultiRange;
+import com.edgedb.driver.datatypes.Range;
+import com.edgedb.driver.exceptions.EdgeDBException;
+import org.jetbrains.annotations.Nullable;
+
+import javax.naming.OperationNotSupportedException;
+import java.util.UUID;
+
+public final class MultiRangeCodec<T> extends CodecBase<MultiRange<T>> {
+    private final RangeCodec<T> rangeCodec;
+
+    @SuppressWarnings("unchecked")
+    public MultiRangeCodec(UUID id, @Nullable CodecMetadata metadata, Class<?> cls, Codec<T> elementCodec) {
+        super(id, metadata, (Class<MultiRange<T>>)cls);
+        this.rangeCodec = new RangeCodec<>(id, metadata, cls, elementCodec);
+    }
+
+    @Override
+    public void serialize(PacketWriter writer, @Nullable MultiRange<T> value, CodecContext context) throws OperationNotSupportedException, EdgeDBException {
+        if(value == null) {
+            return;
+        }
+
+        writer.write(value.length);
+
+        for(int i = 0; i != value.length; i++) {
+            var element = value.get(i);
+            writer.writeDelegateWithLength(w -> rangeCodec.serialize(w, element, context));
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public @Nullable MultiRange<T> deserialize(PacketReader reader, CodecContext context) throws EdgeDBException, OperationNotSupportedException {
+        var length = reader.readInt32();
+
+        if(length == 0) {
+            return MultiRange.empty();
+        }
+
+        var elements = new Range[length];
+
+        for(int i = 0; i != length; i++) {
+            try(var scoped = reader.scopedSlice(reader.readInt32())) {
+                elements[i] = rangeCodec.deserialize(scoped, context);
+            }
+        }
+
+        return new MultiRange<T>(elements);
+    }
+}

src/driver/src/main/java/com/edgedb/driver/binary/protocol/v2/V2ProtocolProvider.java

@@ -42,6 +42,7 @@ public class V2ProtocolProvider extends V1ProtocolProvider implements ProtocolPr
             put(DescriptorType.SCALAR, ScalarTypeDescriptor::new);
             put(DescriptorType.SET, SetTypeDescriptor::new);
             put(DescriptorType.TUPLE, TupleTypeDescriptor::new);
+            put(DescriptorType.MULTI_RANGE, MultiRangeDescriptor::new);
             put(DescriptorType.TYPE_ANNOTATION_TEXT, (ignored, reader) -> new TypeAnnotationTextDescriptor(reader));
         }};
     }
@@ -247,6 +248,22 @@ public TypeDescriptorInfo<DescriptorType> readDescriptor(PacketReader reader) th
                                         t -> Range.empty(t).getClass()
                                 )
                 );
+            case MULTI_RANGE:
+                var multirangeDescriptor = descriptorInfo.as(MultiRangeDescriptor.class);
+
+                return CodecBuilder.getOrCreateCodec(
+                        this,
+                        descriptorInfo.getId(),
+                        metadata,
+                        (id, meta) ->
+                                new CompilableCodec(
+                                        id,
+                                        meta,
+                                        getRelativeCodec.apply(multirangeDescriptor.type.intValue()),
+                                        MultiRangeCodec::new,
+                                        t -> t
+                                )
+                );
             case SCALAR:
                 throw new MissingCodecException(
                         "Could not find the scalar type " + descriptorInfo.getId().toString()

src/driver/src/main/java/com/edgedb/driver/binary/protocol/v2/descriptors/DescriptorType.java

@@ -14,6 +14,7 @@ public enum DescriptorType implements BinaryEnum<Byte> {
     RANGE(9),
     OBJECT(10),
     COMPOUND(11),
+    MULTI_RANGE(12),
     TYPE_ANNOTATION_TEXT(127);
 
     private final byte value;

src/driver/src/main/java/com/edgedb/driver/binary/protocol/v2/descriptors/MultiRangeDescriptor.java

@@ -0,0 +1,43 @@
+package com.edgedb.driver.binary.protocol.v2.descriptors;
+
+import com.edgedb.driver.binary.PacketReader;
+import com.edgedb.driver.binary.codecs.Codec;
+import com.edgedb.driver.binary.protocol.TypeDescriptor;
+import com.edgedb.driver.binary.protocol.TypeDescriptorInfo;
+import com.edgedb.driver.binary.protocol.common.descriptors.CodecMetadata;
+import org.jetbrains.annotations.Nullable;
+import org.joou.UShort;
+
+import java.util.UUID;
+import java.util.function.Function;
+
+public class MultiRangeDescriptor implements TypeDescriptor, MetadataDescriptor {
+
+    public final UUID id;
+    public final String name;
+    public final boolean isSchemaDefined;
+    public final UShort[] ancestors;
+    public final UShort type;
+
+    public MultiRangeDescriptor(UUID id, PacketReader reader) {
+        this.id = id;
+        this.name = reader.readString();
+        this.isSchemaDefined = reader.readBoolean();
+        this.ancestors = reader.readArrayOf(UShort.class, PacketReader::readUInt16, UShort.class);
+        this.type = reader.readUInt16();
+    }
+
+    @Override
+    public UUID getId() {
+        return this.id;
+    }
+
+    @Override
+    public @Nullable CodecMetadata getMetadata(Function<Integer, Codec<?>> getRelativeCodec, Function<Integer, TypeDescriptorInfo<?>> getRelativeDescriptor) {
+        return new CodecMetadata(
+                name,
+                isSchemaDefined,
+                MetadataDescriptor.constructAncestors(ancestors, getRelativeCodec, getRelativeDescriptor)
+        );
+    }
+}

src/driver/src/main/java/com/edgedb/driver/datatypes/MultiRange.java

@@ -0,0 +1,80 @@
+package com.edgedb.driver.datatypes;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
+
+/**
+ * Represents the {@code multirange} type in EdgeDB
+ * @param <T> The inner type of the multirange.
+ */
+@SuppressWarnings("unchecked")
+public class MultiRange<T> {
+    private static final Range<?>[] EMPTY_RANGE_ARRAY = new Range<?>[0];
+    private static final MultiRange<?> EMPTY_MULTI_RANGE = new MultiRange<>();
+
+    /**
+     * Gets the length of this multirange.
+     */
+    public final int length;
+
+    private final Range<T>[] ranges;
+
+    /**
+     * Constructs a new empty multirange
+     */
+    public MultiRange() {
+        ranges = (Range<T>[]) EMPTY_RANGE_ARRAY;
+        length = 0;
+    }
+
+    /**
+     * Constructs a new multirange with the provided elements.
+     * @param elements The elements to construct the multirange with.
+     */
+    public MultiRange(Collection<? extends Range<T>> elements) {
+        ranges = elements.toArray((Range<T>[])EMPTY_RANGE_ARRAY);
+        length = ranges.length;
+    }
+
+    public MultiRange(Range<T>[] elements) {
+        ranges = elements.clone();
+        length = elements.length;
+    }
+
+
+    /**
+     * Gets an element within this multirange by index.
+     * @param i The index of the element to get.
+     * @return The element at the specified index.
+     */
+    public Range<T> get(int i) throws IndexOutOfBoundsException {
+        return ranges[i];
+    }
+
+    /**
+     * Converts this multirange into a hashset.
+     * @return A hashset representing this multirange.
+     */
+    public HashSet<Range<T>> toSet() {
+        return new HashSet<>(Arrays.asList(ranges));
+    }
+
+    public static <U> MultiRange<U> empty(Class<U> cls) {
+        return new MultiRange<>();
+    }
+
+    public static <U> MultiRange<U> empty() {
+        return (MultiRange<U>) EMPTY_MULTI_RANGE;
+    }
+
+    /**
+     * Gets a {@linkplain Class} that represents a multirange of a specified type.
+     * @param cls The inner type of the multirange to represent.
+     * @return A class that represents a multirange of the provided type.
+     * @param <U> The inner type of the multirange.
+     */
+    public static <U> Class<MultiRange<U>> ofType(Class<U> cls) {
+        return (Class<MultiRange<U>>) EMPTY_MULTI_RANGE.getClass();
+    }
+}

src/driver/src/main/java/com/edgedb/driver/datatypes/Range.java

@@ -162,6 +162,17 @@ public static <U> Range<U> empty() {
         return new Range<>(cls, lower, upper, includeLower, includeUpper);
     }
 
+    /**
+     * Gets a {@linkplain Class} that represents a range of a provided type.
+     * @param cls The inner type of the range.
+     * @return A class that represents a range of a provide type.
+     * @param <U> The inner type of the range.
+     */
+    @SuppressWarnings("unchecked")
+    public static <U> Class<Range<U>> ofType(Class<U> cls) {
+        return (Class<Range<U>>) EMPTY_RANGE.getClass();
+    }
+
     /**
      * Gets the element types class of this range.
      * @return The {@linkplain Class<T>} of the element type.

src/driver/src/test/java/QueryTests.java

@@ -1,10 +1,42 @@
 import com.edgedb.driver.EdgeDBClient;
 import com.edgedb.driver.annotations.EdgeDBType;
+import com.edgedb.driver.datatypes.MultiRange;
+import com.edgedb.driver.datatypes.Range;
 import org.junit.jupiter.api.Test;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class QueryTests {
+    @Test
+    public void testMultiRanges() {
+        try(var client = new EdgeDBClient()) {
+            var multiRange = new MultiRange<Long>(new ArrayList<Range<Long>>() {{
+                add(Range.create(Long.class, -40L, -20L));
+                add(Range.create(Long.class, 5L, 10L));
+                add(Range.create(Long.class, 20L, 50L));
+                add(Range.create(Long.class, 5000L, 5001L));
+            }});
+
+            var result = client.queryRequiredSingle(
+                    MultiRange.ofType(Long.class),
+                    "SELECT <multirange<int64>>$arg",
+                    new HashMap<>(){{
+                        put("arg", multiRange);
+                    }}
+            ).toCompletableFuture().get();
+
+            assertThat(result.length).isEqualTo(multiRange.length);
+
+            for(int i = 0; i != multiRange.length; i++) {
+                assertThat(result.get(i)).isEqualTo(multiRange.get(i));
+            }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
 
     @EdgeDBType
     public static class TestDataContainer {
@@ -15,7 +47,7 @@ public static class TestDataContainer {
     }
 
     @Test
-    public void TestPrimitives() {
+    public void testPrimitives() {
         // primitives (long, int, etc.) differ from the class form (Long, Integer, etc.),
         // we test that we can deserialize both in a data structure.
         try(var client = new EdgeDBClient()) {