From a7a8fc33bb38c799cc8db7c0ac86245518b09492 Mon Sep 17 00:00:00 2001
From: Alex Vanin <alexey@nspcc.ru>
Date: Fri, 14 Aug 2020 17:10:15 +0300
Subject: [PATCH] Add stable marshaler for enums

Signed-off-by: Alex Vanin <alexey@nspcc.ru>
---
 util/proto/marshal.go      |   8 +++++
 util/proto/marshal_test.go |  63 ++++++++++++++++++++++++++++++++++++-
 util/proto/test/test.pb.go | Bin 13326 -> 15480 bytes
 util/proto/test/test.proto |   7 +++++
 4 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/util/proto/marshal.go b/util/proto/marshal.go
index 32001766..025d9b45 100644
--- a/util/proto/marshal.go
+++ b/util/proto/marshal.go
@@ -120,6 +120,14 @@ func Int32Size(field int, v int32) int {
 	return UInt64Size(field, uint64(v))
 }
 
+func EnumMarshal(field int, buf []byte, v int32) (int, error) {
+	return UInt64Marshal(field, buf, uint64(v))
+}
+
+func EnumSize(field int, v int32) int {
+	return UInt64Size(field, uint64(v))
+}
+
 // varUIntSize returns length of varint byte sequence for uint64 value 'x'.
 func VarUIntSize(x uint64) int {
 	return (bits.Len64(x|1) + 6) / 7
diff --git a/util/proto/marshal_test.go b/util/proto/marshal_test.go
index 0abcef9c..68025239 100644
--- a/util/proto/marshal_test.go
+++ b/util/proto/marshal_test.go
@@ -10,6 +10,8 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+type SomeEnum int32
+
 type stablePrimitives struct {
 	FieldA []byte
 	FieldB string
@@ -18,8 +20,15 @@ type stablePrimitives struct {
 	FieldE uint32
 	FieldF int64
 	FieldG uint64
+	FieldH SomeEnum
 }
 
+const (
+	ENUM_UNKNOWN  SomeEnum = 0
+	ENUM_POSITIVE          = 1
+	ENUM_NEGATIVE          = -1
+)
+
 func (s *stablePrimitives) stableMarshal(buf []byte, wrongField bool) ([]byte, error) {
 	if s == nil {
 		return []byte{}, nil
@@ -103,6 +112,16 @@ func (s *stablePrimitives) stableMarshal(buf []byte, wrongField bool) ([]byte, e
 	}
 	i += offset
 
+	fieldNum = 300
+	if wrongField {
+		fieldNum++
+	}
+	offset, err = proto.EnumMarshal(fieldNum, buf, int32(s.FieldH))
+	if err != nil {
+		return nil, errors.Wrap(err, "can't marshal field g")
+	}
+	i += offset
+
 	return buf, nil
 }
 
@@ -113,7 +132,8 @@ func (s *stablePrimitives) stableSize() int {
 		proto.Int32Size(201, s.FieldD) +
 		proto.UInt32Size(202, s.FieldE) +
 		proto.Int64Size(203, s.FieldF) +
-		proto.UInt64Size(204, s.FieldG)
+		proto.UInt64Size(204, s.FieldG) +
+		proto.EnumSize(300, int32(s.FieldH))
 }
 
 func TestBytesMarshal(t *testing.T) {
@@ -209,6 +229,14 @@ func TestUInt64Marshal(t *testing.T) {
 	})
 }
 
+func TestEnumMarshal(t *testing.T) {
+	testEnumMarshal(t, ENUM_UNKNOWN, false)
+	testEnumMarshal(t, ENUM_POSITIVE, false)
+	testEnumMarshal(t, ENUM_POSITIVE, true)
+	testEnumMarshal(t, ENUM_NEGATIVE, false)
+	testEnumMarshal(t, ENUM_NEGATIVE, true)
+}
+
 func testBytesMarshal(t *testing.T, data []byte, wrongField bool) {
 	var (
 		wire []byte
@@ -398,3 +426,36 @@ func testUInt64Marshal(t *testing.T, n uint64, wrongField bool) {
 		require.EqualValues(t, 0, result.FieldG)
 	}
 }
+
+func testEnumMarshal(t *testing.T, e SomeEnum, wrongField bool) {
+	var (
+		wire []byte
+		err  error
+
+		custom    = stablePrimitives{FieldH: e}
+		transport = test.Primitives{FieldH: test.Primitives_SomeEnum(e)}
+	)
+
+	wire, err = custom.stableMarshal(nil, wrongField)
+	require.NoError(t, err)
+
+	wireGen, err := transport.Marshal()
+	require.NoError(t, err)
+
+	if !wrongField {
+		// we can check equality because single field cannot be unstable marshalled
+		require.Equal(t, wireGen, wire)
+	} else {
+		require.NotEqual(t, wireGen, wire)
+	}
+
+	result := new(test.Primitives)
+	err = result.Unmarshal(wire)
+	require.NoError(t, err)
+
+	if !wrongField {
+		require.EqualValues(t, custom.FieldH, result.FieldH)
+	} else {
+		require.EqualValues(t, 0, result.FieldH)
+	}
+}
diff --git a/util/proto/test/test.pb.go b/util/proto/test/test.pb.go
index 67908540c2ee967d5de4b5b9ab021df99b49540f..0c3146791ee92bf42e98864599b952a72e7f47fb 100644
GIT binary patch
delta 1705
zcmaJ=&rcIU6b>ydbU{L6tAIv^m4vQr$!@nTZDRu*DoaQTifD`>rQPmI3Eg75B_&ul
z+znCZ;>m+IPm=T>FrMY+f8amBO^?pJ+0qh<dzhVhKfdpK@6DX-{B4|mb=kR%N?FgC
z^)f1%mJ;DOtzy<J6ykjA)l#`+l@9b;ZZo@@U3;EIsLOVdmqya1;PthQm8UDuGS>vl
zmxAqVW_js9(P5brIPM@{?S+;z^JN{SQ8~Z=64Y<itZK=8b;S9k6hgcsn-@8sOwVm)
z)DR{XKAW)_JFzscr#Z;))!RIbVGHO8&jkm;mk4Z!BZXvoZr8`BP?#79i%~Z<6l$=7
zY}Vj*GX=B(6@&!Dh!EkZ>Q=pKqWx;cs?5<*Wl7ng9;1MMh-{Fdh9Gp<dISrILg70D
zYhDpaO|FSQxl)73`S7AUfPcHTyA76?Iq^65VP`vzANFr|A5SkIr#)RD$DSAVqmsT?
zRNF+(zlT7y4HXyaMv5mFfOK3`BuT97TOyr%+MqbO9dWl-G1);CAwv;$Qd!#4Yu4OV
z+|e?RXiAEA0Ps(b$aUo4Kc3l6)sx=vfa7ba-X~tXFg$i!CF8jYT*i-w_(_%XH7`ci
z73ADS{9$-%lQurYU!$+gi@m^|*>)ke?@SJ^r+~i<|8hrS(bnM&Pv1x(N`~4<Ch3|;
zw%(1DTy}&Wv^ZTA$#IjH(vUP31&$b9Bj7H~)4d_nRR*Leta_eW5@0uqum`)rKp`up
zQ~Nx9666O5I}X1*WXB8C&!Xax9RY8Q=~e(0kdh!<0v;JG60A{BO;*Z~NYI2NfF@v7
zAfX6wnnAsG5ghEGfT_qZV4nc>QK*W+C4h&(QHBUe(jdr522O?9fydbcXbc+3Kw|*9
z8f79Ol!??>vH=W5$SATxAwvTs17u4eQeutK@cHQY1gS}b)+rF}v^+t;@O}UCG;zlP
zKxdUP=`rAv;1)<xeB!^$%iu9!E_t>Gjq|MY1iW?TUvl-S_;oPuRs{in4+io1t$zF?
zcoWwGe!MU-gpYz^D-)2tUiQB6;)Kwrwst4h{exP*rX%upKTZYky_pOFlBa|C^Me7h
UJDxckQ_-7ruRGs*JNtg{A5{VIrT_o{

delta 307
zcmexS(U-AdCDUYXR?W%wtc8<z1MxOCk;&d{36o9PMJAsD(p!MECp(0H0!V}S?i@~&
z)i~uQ`*UhdZsFA0ypz+9ZE}Nv#N@vMN0<%GEjAw!WM$+tH_=frsIV{qk(P#=uM1@|
z>7{}=21X##$Q(o(rs{BU8dM}(f;h&8Akq}9$lO9lVY0Qjdmc#IEEyzW1~$qt4XiH_
zL>d@_Nw9eaU|9&6237+hlfimZ!G@Zeg4A1p)k9=~>OmI3HBWYsRM^ZfA;vWMtOECB
XP8prem!%BZHajS~2yRv}3E=<$f`L^6

diff --git a/util/proto/test/test.proto b/util/proto/test/test.proto
index d3b43e64..5aed2b2a 100644
--- a/util/proto/test/test.proto
+++ b/util/proto/test/test.proto
@@ -10,4 +10,11 @@ message Primitives {
     uint32 field_e = 202;
     int64 field_f  = 203;
     uint64 field_g = 204;
+
+    enum SomeEnum {
+        UNKNOWN = 0;
+        POSITIVE = 1;
+        NEGATIVE = -1;
+    }
+    SomeEnum field_h = 300;
 }
\ No newline at end of file