From 09499a01556d68ee740e6c44677d80b09d9c7d7a Mon Sep 17 00:00:00 2001
From: jordy25519 <beauchjord@gmail.com>
Date: Thu, 20 Jul 2023 04:45:38 +1200
Subject: [PATCH] improve rust performance (#944)

* rust specialize get_u8_at & put_u8_at improve performance
use copy_from_slice instead of split_at when r/w u8 array

* remove pointless comment

* get_bytes_at crate scoped

* fix java style lints
---
 .../sbe/generation/rust/LibRsDef.java         | 23 ++++++++++++++-----
 .../sbe/generation/rust/RustGenerator.java    | 15 ++++++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/LibRsDef.java b/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/LibRsDef.java
index 349bc79c9..3d38b28aa 100644
--- a/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/LibRsDef.java
+++ b/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/LibRsDef.java
@@ -175,14 +175,21 @@ static void generateReadBuf(final Appendable writer, final ByteOrder byteOrder)
         indent(writer, 1, "}\n\n");
 
         indent(writer, 1, "#[inline]\n");
-        indent(writer, 1, "fn get_bytes_at<const COUNT: usize>(slice: &[u8], index: usize) -> [u8; COUNT] {\n");
-        indent(writer, 2, "Self::get_bytes(slice.split_at(index).1.split_at(COUNT).0)\n");
+        indent(writer, 1, "pub(crate) fn get_bytes_at<const N: usize>(slice: &[u8], index: usize) -> [u8; N] {\n");
+        indent(writer, 2, "slice[index..index+N].try_into().expect(\"slice with incorrect length\")\n");
         indent(writer, 1, "}\n");
 
         final LinkedHashSet<String> uniquePrimitiveTypes
             = new LinkedHashSet<>(TYPE_NAME_BY_PRIMITIVE_TYPE_MAP.values());
         final String endianness = byteOrder == LITTLE_ENDIAN ? "le" : "be";
 
+        uniquePrimitiveTypes.remove("u8");
+        indent(writer, 0, "\n");
+        indent(writer, 1, "#[inline]\n");
+        indent(writer, 1, "pub fn get_u8_at(&self, index: usize) -> u8 {\n");
+        indent(writer, 2, "self.data[index]\n");
+        indent(writer, 1, "}\n");
+
         for (final String primitiveType : uniquePrimitiveTypes)
         {
             // get_<primitive>_at
@@ -197,7 +204,7 @@ static void generateReadBuf(final Appendable writer, final ByteOrder byteOrder)
         indent(writer, 0, "\n");
         indent(writer, 1, "#[inline]\n");
         indent(writer, 1, "pub fn get_slice_at(&self, index: usize, len: usize) -> &[u8] {\n");
-        indent(writer, 2, "self.data.split_at(index).1.split_at(len).0\n");
+        indent(writer, 2, "&self.data[index..index+len]\n");
         indent(writer, 1, "}\n\n");
 
         writer.append("}\n");
@@ -219,9 +226,7 @@ static void generateWriteBuf(final Writer writer, final ByteOrder byteOrder) thr
         indent(writer, 1, "#[inline]\n");
         indent(writer, 1,
             "pub fn put_bytes_at<const COUNT: usize>(&mut self, index: usize, bytes: [u8; COUNT]) -> usize {\n");
-        indent(writer, 2, "for (i, byte) in bytes.iter().enumerate() {\n");
-        indent(writer, 3, "self.data[index + i] = *byte;\n");
-        indent(writer, 2, "}\n");
+        indent(writer, 2, "self.data[index..index + COUNT].copy_from_slice(&bytes);\n");
         indent(writer, 2, "COUNT\n");
         indent(writer, 1, "}\n\n");
 
@@ -229,6 +234,12 @@ static void generateWriteBuf(final Writer writer, final ByteOrder byteOrder) thr
             = new LinkedHashSet<>(TYPE_NAME_BY_PRIMITIVE_TYPE_MAP.values());
         final String endianness = byteOrder == LITTLE_ENDIAN ? "le" : "be";
 
+        uniquePrimitiveTypes.remove("u8");
+        indent(writer, 1, "#[inline]\n");
+        indent(writer, 1, "pub fn put_u8_at(&mut self, index: usize, value: u8) {\n");
+        indent(writer, 2, "self.data[index] = value;\n");
+        indent(writer, 1, "}\n\n");
+
         for (final String primitiveType : uniquePrimitiveTypes)
         {
             // put_<primitive>_at
diff --git a/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/RustGenerator.java b/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/RustGenerator.java
index bff76d43a..25e200a26 100644
--- a/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/RustGenerator.java
+++ b/sbe-tool/src/main/java/uk/co/real_logic/sbe/generation/rust/RustGenerator.java
@@ -376,6 +376,13 @@ private static void generatePrimitiveEncoder(
             indent(sb, level + 1, "let offset = self.%s;\n", getBufOffset(typeToken));
             indent(sb, level + 1, "let buf = self.get_buf_mut();\n");
 
+            if (rustPrimitiveType.equals("u8"))
+            {
+                indent(sb, level + 1, "buf.put_bytes_at(offset, value);\n");
+                indent(sb, level, "}\n\n");
+                return;
+            }
+
             for (int i = 0; i < arrayLength; i++)
             {
                 if (i == 0)
@@ -646,6 +653,14 @@ private static void generatePrimitiveArrayDecoder(
         }
 
         indent(sb, level + 1, "let buf = self.get_buf();\n");
+        if (rustPrimitiveType.equals("u8"))
+        {
+            indent(sb, level + 1, "ReadBuf::get_bytes_at(buf.data, self.%s)\n",
+                getBufOffset(typeToken));
+            indent(sb, level, "}\n\n");
+            return;
+        }
+
         indent(sb, level + 1, "[\n");
         for (int i = 0; i < arrayLength; i++)
         {