From c33ff03cc569e2f0ec4a66061e3b2bf8f1c0fe0a Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Fri, 11 Mar 2022 15:24:51 -0700
Subject: [PATCH] const-oid: add `ObjectIdentifier::push_arc`

Impl'd as a `const fn`.
---
 const-oid/src/encoder.rs | 19 ++++++++++++++-----
 const-oid/src/lib.rs     |  9 +++++++++
 const-oid/tests/lib.rs   |  9 +++++++++
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/const-oid/src/encoder.rs b/const-oid/src/encoder.rs
index 2c62eb689..877b6ed29 100644
--- a/const-oid/src/encoder.rs
+++ b/const-oid/src/encoder.rs
@@ -32,7 +32,7 @@ enum State {
 }
 
 impl Encoder {
-    /// Create a new encoder initialized to an empty default state
+    /// Create a new encoder initialized to an empty default state.
     pub(crate) const fn new() -> Self {
         Self {
             state: State::Initial,
@@ -41,7 +41,16 @@ impl Encoder {
         }
     }
 
-    /// Encode an [`Arc`] as base 128 into the internal buffer
+    /// Extend an existing OID.
+    pub(crate) const fn extend(oid: ObjectIdentifier) -> Self {
+        Self {
+            state: State::Body,
+            bytes: oid.bytes,
+            cursor: oid.length as usize,
+        }
+    }
+
+    /// Encode an [`Arc`] as base 128 into the internal buffer.
     pub(crate) const fn arc(mut self, arc: Arc) -> Result<Self> {
         match self.state {
             State::Initial => {
@@ -84,7 +93,7 @@ impl Encoder {
         }
     }
 
-    /// Finish encoding an OID
+    /// Finish encoding an OID.
     pub(crate) const fn finish(self) -> Result<ObjectIdentifier> {
         if self.cursor >= 2 {
             Ok(ObjectIdentifier {
@@ -96,7 +105,7 @@ impl Encoder {
         }
     }
 
-    /// Encode a single byte of a base128 value
+    /// Encode a single byte of a Base 128 value.
     const fn encode_base128_byte(mut self, mut n: u32, i: usize, continued: bool) -> Result<Self> {
         let mask = if continued { 0b10000000 } else { 0 };
 
@@ -116,7 +125,7 @@ impl Encoder {
     }
 }
 
-/// Compute the length - 1 of an arc when encoded in base 128
+/// Compute the length - 1 of an arc when encoded in base 128.
 const fn base128_len(arc: Arc) -> usize {
     match arc {
         0..=0x7f => 0,
diff --git a/const-oid/src/lib.rs b/const-oid/src/lib.rs
index 62258cc08..14bf22e0b 100644
--- a/const-oid/src/lib.rs
+++ b/const-oid/src/lib.rs
@@ -157,6 +157,15 @@ impl ObjectIdentifier {
     pub fn arcs(&self) -> Arcs<'_> {
         Arcs::new(self)
     }
+
+    /// Push an additional arc onto this OID, returning the child OID.
+    pub const fn push_arc(self, arc: Arc) -> Result<Self> {
+        // TODO(tarcieri): use `?` when stable in `const fn`
+        match Encoder::extend(self).arc(arc) {
+            Ok(encoder) => encoder.finish(),
+            Err(err) => Err(err),
+        }
+    }
 }
 
 impl AsRef<[u8]> for ObjectIdentifier {
diff --git a/const-oid/tests/lib.rs b/const-oid/tests/lib.rs
index f099263f9..607634f18 100644
--- a/const-oid/tests/lib.rs
+++ b/const-oid/tests/lib.rs
@@ -166,3 +166,12 @@ fn parse_invalid_second_arc() {
         Err(Error::ArcInvalid { arc: 40 })
     );
 }
+
+#[test]
+fn push_arc() {
+    let oid = ObjectIdentifier::new_unwrap("1.2.3");
+    assert_eq!(
+        oid.push_arc(4).unwrap(),
+        ObjectIdentifier::new_unwrap("1.2.3.4")
+    );
+}