Merge pull request protocolbuffers#8473 from Jensaarai/ReadOnlySpan

C#: Add ParseFrom/MergeFrom ReadOnlySpan<byte>
eikemeier · May 17, 2021 · f0da89d · f0da89d
2 parents 7e95c64 + 1252f3e
commit f0da89d
Show file tree

Hide file tree

Showing 6 changed files with 110 additions and 26 deletions.
diff --git a/csharp/src/Google.Protobuf.Test/CodedInputStreamTest.cs b/csharp/src/Google.Protobuf.Test/CodedInputStreamTest.cs
@@ -161,12 +161,21 @@ private static void AssertReadVarintFailure(InvalidProtocolBufferException expec
 
         private static void AssertReadFromParseContext(ReadOnlySequence<byte> input, ParseContextAssertAction assertAction, bool assertIsAtEnd)
         {
+            // Check as ReadOnlySequence<byte>
             ParseContext.Initialize(input, out ParseContext parseCtx);
             assertAction(ref parseCtx);
             if (assertIsAtEnd)
             {
                 Assert.IsTrue(SegmentedBufferHelper.IsAtEnd(ref parseCtx.buffer, ref parseCtx.state));
             }
+
+            // Check as ReadOnlySpan<byte>
+            ParseContext.Initialize(input.ToArray().AsSpan(), out ParseContext spanParseContext);
+            assertAction(ref spanParseContext);
+            if (assertIsAtEnd)
+            {
+                Assert.IsTrue(SegmentedBufferHelper.IsAtEnd(ref spanParseContext.buffer, ref spanParseContext.state));
+            }
         }
 
         [Test]

diff --git a/csharp/src/Google.Protobuf.Test/MessageParsingHelpers.cs b/csharp/src/Google.Protobuf.Test/MessageParsingHelpers.cs
@@ -41,32 +41,38 @@ public static class MessageParsingHelpers
     {
         public static void AssertReadingMessage<T>(MessageParser<T> parser, byte[] bytes, Action<T> assert) where T : IMessage<T>
         {
-            var parsedStream = parser.ParseFrom(bytes);
+            var parsedMsg = parser.ParseFrom(bytes);
+            assert(parsedMsg);
 
             // Load content as single segment
-            var parsedBuffer = parser.ParseFrom(new ReadOnlySequence<byte>(bytes));
-            assert(parsedBuffer);
+            parsedMsg = parser.ParseFrom(new ReadOnlySequence<byte>(bytes));
+            assert(parsedMsg);
 
             // Load content as multiple segments
-            parsedBuffer = parser.ParseFrom(ReadOnlySequenceFactory.CreateWithContent(bytes));
-            assert(parsedBuffer);
+            parsedMsg = parser.ParseFrom(ReadOnlySequenceFactory.CreateWithContent(bytes));
+            assert(parsedMsg);
 
-            assert(parsedStream);
+            // Load content as ReadOnlySpan
+            parsedMsg = parser.ParseFrom(new ReadOnlySpan<byte>(bytes));
+            assert(parsedMsg);
         }
 
         public static void AssertReadingMessage(MessageParser parser, byte[] bytes, Action<IMessage> assert)
         {
-            var parsedStream = parser.ParseFrom(bytes);
+            var parsedMsg = parser.ParseFrom(bytes);
+            assert(parsedMsg);
 
             // Load content as single segment
-            var parsedBuffer = parser.ParseFrom(new ReadOnlySequence<byte>(bytes));
-            assert(parsedBuffer);
+            parsedMsg = parser.ParseFrom(new ReadOnlySequence<byte>(bytes));
+            assert(parsedMsg);
 
             // Load content as multiple segments
-            parsedBuffer = parser.ParseFrom(ReadOnlySequenceFactory.CreateWithContent(bytes));
-            assert(parsedBuffer);
+            parsedMsg = parser.ParseFrom(ReadOnlySequenceFactory.CreateWithContent(bytes));
+            assert(parsedMsg);
 
-            assert(parsedStream);
+            // Load content as ReadOnlySpan
+            parsedMsg = parser.ParseFrom(new ReadOnlySpan<byte>(bytes));
+            assert(parsedMsg);
         }
 
         public static void AssertReadingMessageThrows<TMessage, TException>(MessageParser<TMessage> parser, byte[] bytes)
@@ -76,6 +82,8 @@ public static void AssertReadingMessageThrows<TMessage, TException>(MessageParse
             Assert.Throws<TException>(() => parser.ParseFrom(bytes));
 
             Assert.Throws<TException>(() => parser.ParseFrom(new ReadOnlySequence<byte>(bytes)));
+
+            Assert.Throws<TException>(() => parser.ParseFrom(new ReadOnlySpan<byte>(bytes)));
         }
 
         public static void AssertRoundtrip<T>(MessageParser<T> parser, T message, Action<T> additionalAssert = null) where T : IMessage<T>
@@ -87,20 +95,24 @@ public static void AssertRoundtrip<T>(MessageParser<T> parser, T message, Action
             message.WriteTo(bufferWriter);
             Assert.AreEqual(bytes, bufferWriter.WrittenSpan.ToArray(), "Both serialization approaches need to result in the same data.");
 
+            var parsedMsg = parser.ParseFrom(bytes);
+            Assert.AreEqual(message, parsedMsg);
+            additionalAssert?.Invoke(parsedMsg);
+
             // Load content as single segment
-            var parsedBuffer = parser.ParseFrom(new ReadOnlySequence<byte>(bytes));
-            Assert.AreEqual(message, parsedBuffer);
-            additionalAssert?.Invoke(parsedBuffer);
+            parsedMsg = parser.ParseFrom(new ReadOnlySequence<byte>(bytes));
+            Assert.AreEqual(message, parsedMsg);
+            additionalAssert?.Invoke(parsedMsg);
 
             // Load content as multiple segments
-            parsedBuffer = parser.ParseFrom(ReadOnlySequenceFactory.CreateWithContent(bytes));
-            Assert.AreEqual(message, parsedBuffer);
-            additionalAssert?.Invoke(parsedBuffer);
-
-            var parsedStream = parser.ParseFrom(bytes);
-
-            Assert.AreEqual(message, parsedStream);
-            additionalAssert?.Invoke(parsedStream);
+            parsedMsg = parser.ParseFrom(ReadOnlySequenceFactory.CreateWithContent(bytes));
+            Assert.AreEqual(message, parsedMsg);
+            additionalAssert?.Invoke(parsedMsg);
+
+            // Load content as ReadOnlySpan
+            parsedMsg = parser.ParseFrom(new ReadOnlySpan<byte>(bytes));
+            Assert.AreEqual(message, parsedMsg);
+            additionalAssert?.Invoke(parsedMsg);
         }
 
         public static void AssertWritingMessage(IMessage message)

diff --git a/csharp/src/Google.Protobuf/CodedInputStream.cs b/csharp/src/Google.Protobuf/CodedInputStream.cs
@@ -435,8 +435,7 @@ public void ReadMessage(IMessage builder)
             // we will need to switch back again to CodedInputStream-based parsing (which involves copying and storing the state) to be able to
             // invoke the legacy MergeFrom(CodedInputStream) method.
             // For now, this inefficiency is fine, considering this is only a backward-compatibility scenario (and regenerating the code fixes it).
-            var span = new ReadOnlySpan<byte>(buffer);
-            ParseContext.Initialize(ref span, ref state, out ParseContext ctx);
+            ParseContext.Initialize(buffer.AsSpan(), ref state, out ParseContext ctx);
             try
             {
                 ParsingPrimitivesMessages.ReadMessage(ref ctx, builder);

diff --git a/csharp/src/Google.Protobuf/MessageExtensions.cs b/csharp/src/Google.Protobuf/MessageExtensions.cs
@@ -79,6 +79,15 @@ public static void MergeFrom(this IMessage message, ByteString data) =>
         public static void MergeFrom(this IMessage message, Stream input) =>
             MergeFrom(message, input, false, null);
 
+        /// <summary>
+        /// Merges data from the given span into an existing message.
+        /// </summary>
+        /// <param name="message">The message to merge the data into.</param>
+        /// <param name="span">Span containing the data to merge, which must be protobuf-encoded binary data.</param>
+        [SecuritySafeCritical]
+        public static void MergeFrom(this IMessage message, ReadOnlySpan<byte> span) =>
+            MergeFrom(message, span, false, null);
+
         /// <summary>
         /// Merges length-delimited data from the given stream into an existing message.
         /// </summary>
@@ -294,6 +303,16 @@ internal static void MergeFrom(this IMessage message, ReadOnlySequence<byte> dat
             ParsingPrimitivesMessages.CheckReadEndOfStreamTag(ref ctx.state);
         }
 
+        [SecuritySafeCritical]
+        internal static void MergeFrom(this IMessage message, ReadOnlySpan<byte> data, bool discardUnknownFields, ExtensionRegistry registry)
+        {
+            ParseContext.Initialize(data, out ParseContext ctx);
+            ctx.DiscardUnknownFields = discardUnknownFields;
+            ctx.ExtensionRegistry = registry;
+            ParsingPrimitivesMessages.ReadRawMessage(ref ctx, message);
+            ParsingPrimitivesMessages.CheckReadEndOfStreamTag(ref ctx.state);
+        }
+
         internal static void MergeDelimitedFrom(this IMessage message, Stream input, bool discardUnknownFields, ExtensionRegistry registry)
         {
             ProtoPreconditions.CheckNotNull(message, "message");

diff --git a/csharp/src/Google.Protobuf/MessageParser.cs b/csharp/src/Google.Protobuf/MessageParser.cs
@@ -128,6 +128,19 @@ public IMessage ParseFrom(ReadOnlySequence<byte> data)
             return message;
         }
 
+        /// <summary>
+        /// Parses a message from the given span.
+        /// </summary>
+        /// <param name="data">The data to parse.</param>
+        /// <returns>The parsed message.</returns>
+        [SecuritySafeCritical]
+        public IMessage ParseFrom(ReadOnlySpan<byte> data)
+        {
+            IMessage message = factory();
+            message.MergeFrom(data, DiscardUnknownFields, Extensions);
+            return message;
+        }
+
         /// <summary>
         /// Parses a length-delimited message from the given stream.
         /// </summary>
@@ -315,6 +328,19 @@ internal MessageParser(Func<T> factory, bool discardUnknownFields, ExtensionRegi
             return message;
         }
 
+        /// <summary>
+        /// Parses a message from the given span.
+        /// </summary>
+        /// <param name="data">The data to parse.</param>
+        /// <returns>The parsed message.</returns>
+        [SecuritySafeCritical]
+        public new T ParseFrom(ReadOnlySpan<byte> data)
+        {
+            T message = factory();
+            message.MergeFrom(data, DiscardUnknownFields, Extensions);
+            return message;
+        }
+
         /// <summary>
         /// Parses a length-delimited message from the given stream.
         /// </summary>

diff --git a/csharp/src/Google.Protobuf/ParseContext.cs b/csharp/src/Google.Protobuf/ParseContext.cs
@@ -58,8 +58,27 @@ public ref struct ParseContext
         internal ReadOnlySpan<byte> buffer;
         internal ParserInternalState state;
 
+        /// <summary>
+        /// Initialize a <see cref="ParseContext"/>, building all <see cref="ParserInternalState"/> from defaults and
+        /// the given <paramref name="buffer"/>.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static void Initialize(ReadOnlySpan<byte> buffer, out ParseContext ctx)
+        {
+            ParserInternalState state = default;
+            state.sizeLimit = DefaultSizeLimit;
+            state.recursionLimit = DefaultRecursionLimit;
+            state.currentLimit = int.MaxValue;
+            state.bufferSize = buffer.Length;
+
+            Initialize(buffer, ref state, out ctx);
+        }
+
+        /// <summary>
+        /// Initialize a <see cref="ParseContext"/> using existing <see cref="ParserInternalState"/>, e.g. from <see cref="CodedInputStream"/>.
+        /// </summary>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        internal static void Initialize(ref ReadOnlySpan<byte> buffer, ref ParserInternalState state, out ParseContext ctx)
+        internal static void Initialize(ReadOnlySpan<byte> buffer, ref ParserInternalState state, out ParseContext ctx)
         {
             ctx.buffer = buffer;
             ctx.state = state;