Adding deque utils helper

Author: prwhelan

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/DequeUtils.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.core.inference;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.Writeable;
+
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Deque;
+
+public final class DequeUtils {
+
+    private DequeUtils() {
+        // util functions only
+    }
+
+    public static <T> Deque<T> readDeque(StreamInput in, Writeable.Reader<T> reader) throws IOException {
+        return in.readCollection(ArrayDeque::new, ((stream, deque) -> deque.offer(reader.read(in))));
+    }
+
+    public static boolean dequeEquals(Deque<?> thisDeque, Deque<?> otherDeque) {
+        if (thisDeque.size() != otherDeque.size()) {
+            return false;
+        }
+        var thisIter = thisDeque.iterator();
+        var otherIter = otherDeque.iterator();
+        while (thisIter.hasNext() && otherIter.hasNext()) {
+            if (thisIter.next().equals(otherIter.next()) == false) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    public static int dequeHashCode(Deque<?> deque) {
+        if (deque == null) {
+            return 0;
+        }
+        return deque.stream().reduce(1, (hashCode, chunk) -> 31 * hashCode + (chunk == null ? 0 : chunk.hashCode()), Integer::sum);
+    }
+
+    public static <T> Deque<T> of(T elem) {
+        var deque = new ArrayDeque<T>(1);
+        deque.offer(elem);
+        return deque;
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/StreamingChatCompletionResults.java

@@ -17,11 +17,13 @@
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
-import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.Iterator;
 import java.util.concurrent.Flow;
 
+import static org.elasticsearch.xpack.core.inference.DequeUtils.dequeEquals;
+import static org.elasticsearch.xpack.core.inference.DequeUtils.dequeHashCode;
+import static org.elasticsearch.xpack.core.inference.DequeUtils.readDeque;
 import static org.elasticsearch.xpack.core.inference.results.ChatCompletionResults.COMPLETION;
 
 /**
@@ -40,11 +42,7 @@ public record Results(Deque<Result> results) implements InferenceServiceResults.
         public static final String NAME = "streaming_chat_completion_results";
 
         public Results(StreamInput in) throws IOException {
-            this(deque(in));
-        }
-
-        private static Deque<Result> deque(StreamInput in) throws IOException {
-            return in.readCollection(ArrayDeque::new, ((stream, deque) -> deque.offer(new Result(stream))));
+            this(readDeque(in, Result::new));
         }
 
         @Override
@@ -75,31 +73,10 @@ public boolean equals(Object o) {
             return dequeEquals(this.results, other.results());
         }
 
-        private static boolean dequeEquals(Deque<?> thisDeque, Deque<?> otherDeque) {
-            if (thisDeque.size() != otherDeque.size()) {
-                return false;
-            }
-            var thisIter = thisDeque.iterator();
-            var otherIter = otherDeque.iterator();
-            while (thisIter.hasNext() && otherIter.hasNext()) {
-                if (thisIter.next().equals(otherIter.next()) == false) {
-                    return false;
-                }
-            }
-            return true;
-        }
-
         @Override
         public int hashCode() {
             return dequeHashCode(results);
         }
-
-        private static int dequeHashCode(Deque<?> deque) {
-            if (deque == null) {
-                return 0;
-            }
-            return deque.stream().reduce(1, (hashCode, chunk) -> 31 * hashCode + (chunk == null ? 0 : chunk.hashCode()), Integer::sum);
-        }
     }
 
     public record Result(String delta) implements ChunkedToXContent, Writeable {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/StreamingUnifiedChatCompletionResults.java

@@ -18,14 +18,16 @@
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
-import java.util.ArrayDeque;
 import java.util.Collections;
 import java.util.Deque;
 import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.Flow;
 
 import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.chunk;
+import static org.elasticsearch.xpack.core.inference.DequeUtils.dequeEquals;
+import static org.elasticsearch.xpack.core.inference.DequeUtils.dequeHashCode;
+import static org.elasticsearch.xpack.core.inference.DequeUtils.readDeque;
 
 /**
  * Chat Completion results that only contain a Flow.Publisher.
@@ -69,11 +71,7 @@ public record Results(Deque<ChatCompletionChunk> chunks) implements InferenceSer
         public static String NAME = "streaming_unified_chat_completion_results";
 
         public Results(StreamInput in) throws IOException {
-            this(deque(in));
-        }
-
-        private static Deque<ChatCompletionChunk> deque(StreamInput in) throws IOException {
-            return in.readCollection(ArrayDeque::new, ((stream, deque) -> deque.offer(new ChatCompletionChunk(stream))));
+            this(readDeque(in, ChatCompletionChunk::new));
         }
 
         @Override
@@ -98,31 +96,10 @@ public boolean equals(Object o) {
             return dequeEquals(chunks, results.chunks());
         }
 
-        private static boolean dequeEquals(Deque<?> thisDeque, Deque<?> otherDeque) {
-            if (thisDeque.size() != otherDeque.size()) {
-                return false;
-            }
-            var thisIter = thisDeque.iterator();
-            var otherIter = otherDeque.iterator();
-            while (thisIter.hasNext() && otherIter.hasNext()) {
-                if (thisIter.next().equals(otherIter.next()) == false) {
-                    return false;
-                }
-            }
-            return true;
-        }
-
         @Override
         public int hashCode() {
             return dequeHashCode(chunks);
         }
-
-        private static int dequeHashCode(Deque<?> deque) {
-            if (deque == null) {
-                return 0;
-            }
-            return deque.stream().reduce(1, (hashCode, chunk) -> 31 * hashCode + (chunk == null ? 0 : chunk.hashCode()), Integer::sum);
-        }
     }
 
     public record ChatCompletionChunk(String id, List<Choice> choices, String model, String object, ChatCompletionChunk.Usage usage)

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/DequeUtilsTests.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.core.inference;
+
+import org.elasticsearch.common.io.stream.ByteArrayStreamInput;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.test.ESTestCase;
+
+import java.io.IOException;
+
+import static org.mockito.Mockito.mock;
+
+public class DequeUtilsTests extends ESTestCase {
+
+    public void testEqualsAndHashCodeWithSameObject() {
+        var someObject = mock();
+        var dequeOne = DequeUtils.of(someObject);
+        var dequeTwo = DequeUtils.of(someObject);
+        assertTrue(DequeUtils.dequeEquals(dequeOne, dequeTwo));
+        assertEquals(DequeUtils.dequeHashCode(dequeOne), DequeUtils.dequeHashCode(dequeTwo));
+    }
+
+    public void testEqualsAndHashCodeWithEqualsObject() {
+        var dequeOne = DequeUtils.of("the same string");
+        var dequeTwo = DequeUtils.of("the same string");
+        assertTrue(DequeUtils.dequeEquals(dequeOne, dequeTwo));
+        assertEquals(DequeUtils.dequeHashCode(dequeOne), DequeUtils.dequeHashCode(dequeTwo));
+    }
+
+    public void testNotEqualsAndHashCode() {
+        var dequeOne = DequeUtils.of(mock());
+        var dequeTwo = DequeUtils.of(mock());
+        assertFalse(DequeUtils.dequeEquals(dequeOne, dequeTwo));
+        assertNotEquals(DequeUtils.dequeHashCode(dequeOne), DequeUtils.dequeHashCode(dequeTwo));
+    }
+
+    public void testReadFromStream() throws IOException {
+        var dequeOne = DequeUtils.of("this is a string");
+        var out = new BytesStreamOutput();
+        out.writeStringCollection(dequeOne);
+        var in = new ByteArrayStreamInput(out.bytes().array());
+        var dequeTwo = DequeUtils.readDeque(in, StreamInput::readString);
+        assertTrue(DequeUtils.dequeEquals(dequeOne, dequeTwo));
+        assertEquals(DequeUtils.dequeHashCode(dequeOne), DequeUtils.dequeHashCode(dequeTwo));
+    }
+}