Add KeyBalls to shot proto

rbproto/python/shot_pb2.py

@@ -13,7 +13,7 @@ _sym_db = _symbol_database.Default()
 
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\nshot.proto\x12\x0fserialized_shot\"?\n\x03\x42ox\x12\x0c\n\x04left\x18\x01 \x01(\x02\x12\x0b\n\x03top\x18\x02 \x01(\x02\x12\r\n\x05width\x18\x03 \x01(\x02\x12\x0e\n\x06height\x18\x04 \x01(\x02\"\x1d\n\x05Point\x12\t\n\x01x\x18\x01 \x01(\x02\x12\t\n\x01y\x18\x02 \x01(\x02\"\x7f\n\rBallDetection\x12.\n\x0eplane_position\x18\x01 \x01(\x0b\x32\x16.serialized_shot.Point\x12(\n\nannotation\x18\x02 \x01(\x0b\x32\x14.serialized_shot.Box\x12\x14\n\x0cinterpolated\x18\x03 \x01(\x08\"T\n\x10RLEBallDetection\x12\x31\n\tdetection\x18\x01 \x01(\x0b\x32\x1e.serialized_shot.BallDetection\x12\r\n\x05\x63ount\x18\x02 \x01(\r\"L\n\x13RLEDetectionHistory\x12\x35\n\ndetections\x18\x01 \x03(\x0b\x32!.serialized_shot.RLEBallDetection\"F\n\x10\x44\x65tectionHistory\x12\x32\n\ndetections\x18\x01 \x03(\x0b\x32\x1e.serialized_shot.BallDetection\"\xfc\x01\n\rCollisionInfo\x12\x0e\n\x06source\x18\x01 \x01(\r\x12M\n\x10\x62\x61ll_identifiers\x18\x02 \x03(\x0b\x32\x33.serialized_shot.CollisionInfo.BallIdentifiersEntry\x12\x17\n\x0fwall_identifier\x18\x03 \x01(\r\x12\x13\n\x0b\x66rame_index\x18\x04 \x01(\r\x12\x0e\n\x06static\x18\x05 \x01(\x08\x1aN\n\x14\x42\x61llIdentifiersEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.serialized_shot.Point:\x02\x38\x01\"\xcc\x02\n\x04Path\x12\x13\n\x0bstart_frame\x18\x01 \x01(\r\x12\x11\n\tend_frame\x18\x02 \x01(\r\x12\x37\n\ndetections\x18\x03 \x01(\x0b\x32!.serialized_shot.DetectionHistoryH\x00\x12>\n\x0erle_detections\x18\x04 \x01(\x0b\x32$.serialized_shot.RLEDetectionHistoryH\x00\x12\x15\n\x0bnot_present\x18\x05 \x01(\x08H\x00\x12\x11\n\tis_static\x18\x06 \x01(\x08\x12\x32\n\nstart_info\x18\x07 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfo\x12\x30\n\x08\x65nd_info\x18\x08 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfoB\x13\n\x11\x64\x65tection_history\"R\n\x11IdentifierHistory\x12\x17\n\x0f\x62\x61ll_identifier\x18\x01 \x01(\r\x12$\n\x05paths\x18\x02 \x03(\x0b\x32\x15.serialized_shot.Path\"H\n\x04Shot\x12@\n\x14identifier_histories\x18\x03 \x03(\x0b\x32\".serialized_shot.IdentifierHistoryb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\nshot.proto\x12\x0fserialized_shot\"?\n\x03\x42ox\x12\x0c\n\x04left\x18\x01 \x01(\x02\x12\x0b\n\x03top\x18\x02 \x01(\x02\x12\r\n\x05width\x18\x03 \x01(\x02\x12\x0e\n\x06height\x18\x04 \x01(\x02\"\x1d\n\x05Point\x12\t\n\x01x\x18\x01 \x01(\x02\x12\t\n\x01y\x18\x02 \x01(\x02\"\x7f\n\rBallDetection\x12.\n\x0eplane_position\x18\x01 \x01(\x0b\x32\x16.serialized_shot.Point\x12(\n\nannotation\x18\x02 \x01(\x0b\x32\x14.serialized_shot.Box\x12\x14\n\x0cinterpolated\x18\x03 \x01(\x08\"T\n\x10RLEBallDetection\x12\x31\n\tdetection\x18\x01 \x01(\x0b\x32\x1e.serialized_shot.BallDetection\x12\r\n\x05\x63ount\x18\x02 \x01(\r\"L\n\x13RLEDetectionHistory\x12\x35\n\ndetections\x18\x01 \x03(\x0b\x32!.serialized_shot.RLEBallDetection\"F\n\x10\x44\x65tectionHistory\x12\x32\n\ndetections\x18\x01 \x03(\x0b\x32\x1e.serialized_shot.BallDetection\"\xfc\x01\n\rCollisionInfo\x12\x0e\n\x06source\x18\x01 \x01(\r\x12M\n\x10\x62\x61ll_identifiers\x18\x02 \x03(\x0b\x32\x33.serialized_shot.CollisionInfo.BallIdentifiersEntry\x12\x17\n\x0fwall_identifier\x18\x03 \x01(\r\x12\x13\n\x0b\x66rame_index\x18\x04 \x01(\r\x12\x0e\n\x06static\x18\x05 \x01(\x08\x1aN\n\x14\x42\x61llIdentifiersEntry\x12\x0b\n\x03key\x18\x01 \x01(\r\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.serialized_shot.Point:\x02\x38\x01\"\xcc\x02\n\x04Path\x12\x13\n\x0bstart_frame\x18\x01 \x01(\r\x12\x11\n\tend_frame\x18\x02 \x01(\r\x12\x37\n\ndetections\x18\x03 \x01(\x0b\x32!.serialized_shot.DetectionHistoryH\x00\x12>\n\x0erle_detections\x18\x04 \x01(\x0b\x32$.serialized_shot.RLEDetectionHistoryH\x00\x12\x15\n\x0bnot_present\x18\x05 \x01(\x08H\x00\x12\x11\n\tis_static\x18\x06 \x01(\x08\x12\x32\n\nstart_info\x18\x07 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfo\x12\x30\n\x08\x65nd_info\x18\x08 \x01(\x0b\x32\x1e.serialized_shot.CollisionInfoB\x13\n\x11\x64\x65tection_history\"R\n\x11IdentifierHistory\x12\x17\n\x0f\x62\x61ll_identifier\x18\x01 \x01(\r\x12$\n\x05paths\x18\x02 \x03(\x0b\x32\x15.serialized_shot.Path\"j\n\x12KeyBallIdentifiers\x12\x10\n\x08\x63ue_ball\x18\x01 \x01(\r\x12\x13\n\x0bobject_ball\x18\x02 \x01(\r\x12\x13\n\x0btarget_ball\x18\x03 \x01(\r\x12\x18\n\x10\x63ontact_sequence\x18\x04 \x03(\r\"\x80\x01\n\x04Shot\x12@\n\x14identifier_histories\x18\x03 \x03(\x0b\x32\".serialized_shot.IdentifierHistory\x12\x36\n\tkey_balls\x18\x04 \x01(\x0b\x32#.serialized_shot.KeyBallIdentifiersb\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -42,6 +42,8 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_PATH']._serialized_end=1080
   _globals['_IDENTIFIERHISTORY']._serialized_start=1082
   _globals['_IDENTIFIERHISTORY']._serialized_end=1164
-  _globals['_SHOT']._serialized_start=1166
+  _globals['_KEYBALLIDENTIFIERS']._serialized_start=1166
-  _globals['_SHOT']._serialized_end=1238
+  _globals['_KEYBALLIDENTIFIERS']._serialized_end=1272
+  _globals['_SHOT']._serialized_start=1275
+  _globals['_SHOT']._serialized_end=1403
 # @@protoc_insertion_point(module_scope)

rbproto/python/shot_pb2.pyi

@@ -104,8 +104,22 @@ class IdentifierHistory(_message.Message):
     paths: _containers.RepeatedCompositeFieldContainer[Path]
     def __init__(self, ball_identifier: _Optional[int] = ..., paths: _Optional[_Iterable[_Union[Path, _Mapping]]] = ...) -> None: ...
 
+class KeyBallIdentifiers(_message.Message):
+    __slots__ = ["cue_ball", "object_ball", "target_ball", "contact_sequence"]
+    CUE_BALL_FIELD_NUMBER: _ClassVar[int]
+    OBJECT_BALL_FIELD_NUMBER: _ClassVar[int]
+    TARGET_BALL_FIELD_NUMBER: _ClassVar[int]
+    CONTACT_SEQUENCE_FIELD_NUMBER: _ClassVar[int]
+    cue_ball: int
+    object_ball: int
+    target_ball: int
+    contact_sequence: _containers.RepeatedScalarFieldContainer[int]
+    def __init__(self, cue_ball: _Optional[int] = ..., object_ball: _Optional[int] = ..., target_ball: _Optional[int] = ..., contact_sequence: _Optional[_Iterable[int]] = ...) -> None: ...
+
 class Shot(_message.Message):
-    __slots__ = ["identifier_histories"]
+    __slots__ = ["identifier_histories", "key_balls"]
     IDENTIFIER_HISTORIES_FIELD_NUMBER: _ClassVar[int]
+    KEY_BALLS_FIELD_NUMBER: _ClassVar[int]
     identifier_histories: _containers.RepeatedCompositeFieldContainer[IdentifierHistory]
-    def __init__(self, identifier_histories: _Optional[_Iterable[_Union[IdentifierHistory, _Mapping]]] = ...) -> None: ...
+    key_balls: KeyBallIdentifiers
+    def __init__(self, identifier_histories: _Optional[_Iterable[_Union[IdentifierHistory, _Mapping]]] = ..., key_balls: _Optional[_Union[KeyBallIdentifiers, _Mapping]] = ...) -> None: ...

rbproto/shot.proto

@@ -58,7 +58,18 @@ message IdentifierHistory {
   uint32 ball_identifier = 1;
   repeated Path paths = 2;
 }
+message KeyBallIdentifiers {
+  uint32 cue_ball = 1;
+  uint32 object_ball = 2;
+  uint32 target_ball = 3;
+
+  // For now this will just be cue->object/target->target
+  // Long term this will potentially represent a linked list
+  // of all balls in a shot.
+  repeated uint32 contact_sequence = 4;
+}
 
 message Shot {
   repeated IdentifierHistory identifier_histories = 3;
+  KeyBallIdentifiers key_balls = 4;
 }

rbproto/ts/shot.ts

@@ -815,8 +815,121 @@ function _decodeIdentifierHistory(bb: ByteBuffer): IdentifierHistory {
   return message;
 }
 
+export interface KeyBallIdentifiers {
+  cue_ball?: number;
+  object_ball?: number;
+  target_ball?: number;
+  contact_sequence?: number[];
+}
+
+export function encodeKeyBallIdentifiers(
+  message: KeyBallIdentifiers,
+): Uint8Array {
+  let bb = popByteBuffer();
+  _encodeKeyBallIdentifiers(message, bb);
+  return toUint8Array(bb);
+}
+
+function _encodeKeyBallIdentifiers(
+  message: KeyBallIdentifiers,
+  bb: ByteBuffer,
+): void {
+  // optional uint32 cue_ball = 1;
+  let $cue_ball = message.cue_ball;
+  if ($cue_ball !== undefined) {
+    writeVarint32(bb, 8);
+    writeVarint32(bb, $cue_ball);
+  }
+
+  // optional uint32 object_ball = 2;
+  let $object_ball = message.object_ball;
+  if ($object_ball !== undefined) {
+    writeVarint32(bb, 16);
+    writeVarint32(bb, $object_ball);
+  }
+
+  // optional uint32 target_ball = 3;
+  let $target_ball = message.target_ball;
+  if ($target_ball !== undefined) {
+    writeVarint32(bb, 24);
+    writeVarint32(bb, $target_ball);
+  }
+
+  // repeated uint32 contact_sequence = 4;
+  let array$contact_sequence = message.contact_sequence;
+  if (array$contact_sequence !== undefined) {
+    let packed = popByteBuffer();
+    for (let value of array$contact_sequence) {
+      writeVarint32(packed, value);
+    }
+    writeVarint32(bb, 34);
+    writeVarint32(bb, packed.offset);
+    writeByteBuffer(bb, packed);
+    pushByteBuffer(packed);
+  }
+}
+
+export function decodeKeyBallIdentifiers(
+  binary: Uint8Array,
+): KeyBallIdentifiers {
+  return _decodeKeyBallIdentifiers(wrapByteBuffer(binary));
+}
+
+function _decodeKeyBallIdentifiers(bb: ByteBuffer): KeyBallIdentifiers {
+  let message: KeyBallIdentifiers = {} as any;
+
+  end_of_message: while (!isAtEnd(bb)) {
+    let tag = readVarint32(bb);
+
+    switch (tag >>> 3) {
+      case 0:
+        break end_of_message;
+
+      // optional uint32 cue_ball = 1;
+      case 1: {
+        message.cue_ball = readVarint32(bb) >>> 0;
+        break;
+      }
+
+      // optional uint32 object_ball = 2;
+      case 2: {
+        message.object_ball = readVarint32(bb) >>> 0;
+        break;
+      }
+
+      // optional uint32 target_ball = 3;
+      case 3: {
+        message.target_ball = readVarint32(bb) >>> 0;
+        break;
+      }
+
+      // repeated uint32 contact_sequence = 4;
+      case 4: {
+        let values =
+          message.contact_sequence || (message.contact_sequence = []);
+        if ((tag & 7) === 2) {
+          let outerLimit = pushTemporaryLength(bb);
+          while (!isAtEnd(bb)) {
+            values.push(readVarint32(bb) >>> 0);
+          }
+          bb.limit = outerLimit;
+        } else {
+          values.push(readVarint32(bb) >>> 0);
+        }
+        break;
+      }
+
+      default:
+        skipUnknownField(bb, tag & 7);
+    }
+  }
+
+  return message;
+}
+
 export interface Shot {
   identifier_histories?: IdentifierHistory[];
+  key_balls?: KeyBallIdentifiers;
 }
 
 export function encodeShot(message: Shot): Uint8Array {
@@ -838,6 +951,17 @@ function _encodeShot(message: Shot, bb: ByteBuffer): void {
       pushByteBuffer(nested);
     }
   }
+
+  // optional KeyBallIdentifiers key_balls = 4;
+  let $key_balls = message.key_balls;
+  if ($key_balls !== undefined) {
+    writeVarint32(bb, 34);
+    let nested = popByteBuffer();
+    _encodeKeyBallIdentifiers($key_balls, nested);
+    writeVarint32(bb, nested.limit);
+    writeByteBuffer(bb, nested);
+    pushByteBuffer(nested);
+  }
 }
 
 export function decodeShot(binary: Uint8Array): Shot {
@@ -864,6 +988,14 @@ function _decodeShot(bb: ByteBuffer): Shot {
         break;
       }
 
+      // optional KeyBallIdentifiers key_balls = 4;
+      case 4: {
+        let limit = pushTemporaryLength(bb);
+        message.key_balls = _decodeKeyBallIdentifiers(bb);
+        bb.limit = limit;
+        break;
+      }
+
       default:
         skipUnknownField(bb, tag & 7);
     }