22
33import com .google .protobuf .InvalidProtocolBufferException ;
44import io .github .dfa1 .vortex .proto .EncodingProtos ;
5+ import io .github .dfa1 .vortex .proto .ScalarProtos ;
56import io .github .dfa1 .vortex .core .array .Array ;
67import io .github .dfa1 .vortex .core .ArrayStats ;
78import io .github .dfa1 .vortex .core .DType ;
1819import java .lang .foreign .MemorySegment ;
1920import java .lang .foreign .ValueLayout ;
2021import java .nio .ByteBuffer ;
22+ import java .util .List ;
2123
2224/// Decoder for {@code vortex.sequence}: {@code A[i] = base + i * multiplier}.
2325///
@@ -33,14 +35,112 @@ public EncodingId encodingId() {
3335
3436 @ Override
3537 public EncodeResult encode (DType dtype , Object data ) {
36- throw new UnsupportedOperationException ( " encode not supported by " + encodingId () );
38+ return Encoder . encode ( dtype , data );
3739 }
3840
3941 @ Override
4042 public Array decode (DecodeContext ctx ) {
4143 return Decoder .decode (ctx );
4244 }
4345
46+ private static final class Encoder {
47+
48+ private static EncodeResult encode (DType dtype , Object data ) {
49+ if (!(dtype instanceof DType .Primitive p )) {
50+ throw new VortexException (EncodingId .VORTEX_SEQUENCE , "encode only supports Primitive dtype, got " + dtype );
51+ }
52+ PType pt = p .ptype ();
53+ return switch (pt ) {
54+ case I8 , I16 , I32 , I64 , U8 , U16 , U32 , U64 -> encodeInteger (pt , data );
55+ case F32 -> encodeF32 ((float []) data );
56+ case F64 -> encodeF64 ((double []) data );
57+ case F16 -> throw new VortexException (EncodingId .VORTEX_SEQUENCE , "F16 not supported" );
58+ };
59+ }
60+
61+ private static EncodeResult encodeInteger (PType pt , Object data ) {
62+ int n = intArrayLength (pt , data );
63+ long base = 0 ;
64+ long multiplier = 0 ;
65+ if (n > 0 ) {
66+ base = readLong (pt , data , 0 );
67+ multiplier = n > 1 ? readLong (pt , data , 1 ) - base : 0 ;
68+ for (int i = 2 ; i < n ; i ++) {
69+ long expected = base + (long ) i * multiplier ;
70+ if (readLong (pt , data , i ) != expected ) {
71+ throw new VortexException (EncodingId .VORTEX_SEQUENCE , "not an arithmetic sequence at index " + i );
72+ }
73+ }
74+ }
75+ ScalarProtos .ScalarValue baseScalar = buildIntScalar (pt , base );
76+ ScalarProtos .ScalarValue mulScalar = buildIntScalar (pt , multiplier );
77+ return buildResult (baseScalar , mulScalar );
78+ }
79+
80+ private static EncodeResult encodeF32 (float [] data ) {
81+ float base = data .length > 0 ? data [0 ] : 0f ;
82+ float mul = data .length > 1 ? data [1 ] - base : 0f ;
83+ for (int i = 2 ; i < data .length ; i ++) {
84+ if (data [i ] != base + (float ) i * mul ) {
85+ throw new VortexException (EncodingId .VORTEX_SEQUENCE , "not an arithmetic sequence at index " + i );
86+ }
87+ }
88+ ScalarProtos .ScalarValue baseScalar = ScalarProtos .ScalarValue .newBuilder ().setF32Value (base ).build ();
89+ ScalarProtos .ScalarValue mulScalar = ScalarProtos .ScalarValue .newBuilder ().setF32Value (mul ).build ();
90+ return buildResult (baseScalar , mulScalar );
91+ }
92+
93+ private static EncodeResult encodeF64 (double [] data ) {
94+ double base = data .length > 0 ? data [0 ] : 0.0 ;
95+ double mul = data .length > 1 ? data [1 ] - base : 0.0 ;
96+ for (int i = 2 ; i < data .length ; i ++) {
97+ if (data [i ] != base + (double ) i * mul ) {
98+ throw new VortexException (EncodingId .VORTEX_SEQUENCE , "not an arithmetic sequence at index " + i );
99+ }
100+ }
101+ ScalarProtos .ScalarValue baseScalar = ScalarProtos .ScalarValue .newBuilder ().setF64Value (base ).build ();
102+ ScalarProtos .ScalarValue mulScalar = ScalarProtos .ScalarValue .newBuilder ().setF64Value (mul ).build ();
103+ return buildResult (baseScalar , mulScalar );
104+ }
105+
106+ private static EncodeResult buildResult (ScalarProtos .ScalarValue base , ScalarProtos .ScalarValue mul ) {
107+ EncodingProtos .SequenceMetadata meta = EncodingProtos .SequenceMetadata .newBuilder ()
108+ .setBase (base )
109+ .setMultiplier (mul )
110+ .build ();
111+ ByteBuffer metaBuf = ByteBuffer .wrap (meta .toByteArray ());
112+ EncodeNode node = new EncodeNode (EncodingId .VORTEX_SEQUENCE , metaBuf , new EncodeNode [0 ], new int []{});
113+ return new EncodeResult (node , List .of (), null , null );
114+ }
115+
116+ private static ScalarProtos .ScalarValue buildIntScalar (PType pt , long value ) {
117+ return switch (pt ) {
118+ case U8 , U16 , U32 , U64 -> ScalarProtos .ScalarValue .newBuilder ().setUint64Value (value ).build ();
119+ default -> ScalarProtos .ScalarValue .newBuilder ().setInt64Value (value ).build ();
120+ };
121+ }
122+
123+ private static int intArrayLength (PType pt , Object data ) {
124+ return switch (pt ) {
125+ case I8 , U8 -> ((byte []) data ).length ;
126+ case I16 , U16 -> ((short []) data ).length ;
127+ case I32 , U32 -> ((int []) data ).length ;
128+ case I64 , U64 -> ((long []) data ).length ;
129+ default -> throw new VortexException (EncodingId .VORTEX_SEQUENCE , "unsupported ptype: " + pt );
130+ };
131+ }
132+
133+ private static long readLong (PType pt , Object data , int i ) {
134+ return switch (pt ) {
135+ case I8 , U8 -> ((byte []) data )[i ];
136+ case I16 , U16 -> ((short []) data )[i ];
137+ case I32 , U32 -> ((int []) data )[i ];
138+ case I64 , U64 -> ((long []) data )[i ];
139+ default -> throw new VortexException (EncodingId .VORTEX_SEQUENCE , "unsupported ptype: " + pt );
140+ };
141+ }
142+ }
143+
44144 private static final class Decoder {
45145
46146 private static Array decode (DecodeContext ctx ) {
0 commit comments