From 610dc9f9a3f87bff60111db44cc9c9e838bfbfe2 Mon Sep 17 00:00:00 2001
From: Michael Macias <zaeleus@gmail.com>
Date: Mon, 22 Jul 2024 16:59:45 -0500
Subject: [PATCH] sam/alignment/record/sequence: Add checked split at

---
 noodles-bam/src/record/sequence.rs            | 15 +++++++++++++
 .../src/record/sequence/subsequence.rs        | 15 +++++++++++++
 noodles-sam/CHANGELOG.md                      |  2 +-
 noodles-sam/src/alignment/record/sequence.rs  | 21 +++++++++++++++++++
 .../src/alignment/record_buf/sequence.rs      | 18 ++++++++++++++++
 noodles-sam/src/record/sequence.rs            | 20 +++++++++++++++++-
 6 files changed, 89 insertions(+), 2 deletions(-)

diff --git a/noodles-bam/src/record/sequence.rs b/noodles-bam/src/record/sequence.rs
index 13991cfe5..b99f916bb 100644
--- a/noodles-bam/src/record/sequence.rs
+++ b/noodles-bam/src/record/sequence.rs
@@ -82,6 +82,21 @@ impl<'a> sam::alignment::record::Sequence for Sequence<'a> {
         self.get(i)
     }
 
+    fn split_at_checked(
+        &self,
+        mid: usize,
+    ) -> Option<(
+        Box<dyn sam::alignment::record::Sequence + '_>,
+        Box<dyn sam::alignment::record::Sequence + '_>,
+    )> {
+        self.split_at_checked(mid).map(|(left, right)| {
+            (
+                Box::new(left) as Box<dyn sam::alignment::record::Sequence + '_>,
+                Box::new(right) as Box<dyn sam::alignment::record::Sequence + '_>,
+            )
+        })
+    }
+
     fn iter(&self) -> Box<dyn Iterator<Item = u8> + '_> {
         Box::new(self.iter())
     }
diff --git a/noodles-bam/src/record/sequence/subsequence.rs b/noodles-bam/src/record/sequence/subsequence.rs
index eb1d24e6d..083260111 100644
--- a/noodles-bam/src/record/sequence/subsequence.rs
+++ b/noodles-bam/src/record/sequence/subsequence.rs
@@ -77,6 +77,21 @@ impl<'a> sam::alignment::record::Sequence for Subsequence<'a> {
         self.get(i)
     }
 
+    fn split_at_checked(
+        &self,
+        mid: usize,
+    ) -> Option<(
+        Box<dyn sam::alignment::record::Sequence + '_>,
+        Box<dyn sam::alignment::record::Sequence + '_>,
+    )> {
+        self.split_at_checked(mid).map(|(left, right)| {
+            (
+                Box::new(left) as Box<dyn sam::alignment::record::Sequence + '_>,
+                Box::new(right) as Box<dyn sam::alignment::record::Sequence + '_>,
+            )
+        })
+    }
+
     fn iter(&self) -> Box<dyn Iterator<Item = u8> + '_> {
         Box::new(self.iter())
     }
diff --git a/noodles-sam/CHANGELOG.md b/noodles-sam/CHANGELOG.md
index a0f00663f..e2c92640a 100644
--- a/noodles-sam/CHANGELOG.md
+++ b/noodles-sam/CHANGELOG.md
@@ -5,7 +5,7 @@
 ### Added
 
   * sam/alignment/record/sequence: Add get base by index (`Sequence::get`)
-    ([#283]).
+    ([#283]) and checked split at (`Sequence::split_at_checked`).
 
 [#283]: https://github.com/zaeleus/noodles/issues/283
 
diff --git a/noodles-sam/src/alignment/record/sequence.rs b/noodles-sam/src/alignment/record/sequence.rs
index bdb970cd2..2f017ecc5 100644
--- a/noodles-sam/src/alignment/record/sequence.rs
+++ b/noodles-sam/src/alignment/record/sequence.rs
@@ -9,6 +9,12 @@ pub trait Sequence {
     /// Returns the base at the given index.
     fn get(&self, i: usize) -> Option<u8>;
 
+    /// Splits the subsequence into two subsequences at the given index.
+    fn split_at_checked(
+        &self,
+        mid: usize,
+    ) -> Option<(Box<dyn Sequence + '_>, Box<dyn Sequence + '_>)>;
+
     /// Returns an iterator over bases.
     fn iter(&self) -> Box<dyn Iterator<Item = u8> + '_>;
 }
@@ -35,6 +41,13 @@ impl Sequence for Box<dyn Sequence + '_> {
         (**self).get(i)
     }
 
+    fn split_at_checked(
+        &self,
+        mid: usize,
+    ) -> Option<(Box<dyn Sequence + '_>, Box<dyn Sequence + '_>)> {
+        (**self).split_at_checked(mid)
+    }
+
     fn iter(&self) -> Box<dyn Iterator<Item = u8> + '_> {
         (**self).iter()
     }
@@ -61,6 +74,14 @@ mod tests {
                 self.0.get(i).copied()
             }
 
+            fn split_at_checked(
+                &self,
+                mid: usize,
+            ) -> Option<(Box<dyn Sequence + '_>, Box<dyn Sequence + '_>)> {
+                let (left, right) = self.0.split_at(mid);
+                Some((Box::new(T(left.to_vec())), Box::new(T(right.to_vec()))))
+            }
+
             fn iter(&self) -> Box<dyn Iterator<Item = u8> + '_> {
                 Box::new(self.0.iter().copied())
             }
diff --git a/noodles-sam/src/alignment/record_buf/sequence.rs b/noodles-sam/src/alignment/record_buf/sequence.rs
index 352f47287..6ebc97abc 100644
--- a/noodles-sam/src/alignment/record_buf/sequence.rs
+++ b/noodles-sam/src/alignment/record_buf/sequence.rs
@@ -92,6 +92,24 @@ impl crate::alignment::record::Sequence for &Sequence {
         self.0.get(i).copied()
     }
 
+    fn split_at_checked(
+        &self,
+        mid: usize,
+    ) -> Option<(
+        Box<dyn crate::alignment::record::Sequence + '_>,
+        Box<dyn crate::alignment::record::Sequence + '_>,
+    )> {
+        if mid > self.len() {
+            let (left, right) = self.0.split_at(mid);
+            Some((
+                Box::new(crate::record::Sequence::new(left)),
+                Box::new(crate::record::Sequence::new(right)),
+            ))
+        } else {
+            None
+        }
+    }
+
     fn iter(&self) -> Box<dyn Iterator<Item = u8> + '_> {
         Box::new(self.0.iter().copied())
     }
diff --git a/noodles-sam/src/record/sequence.rs b/noodles-sam/src/record/sequence.rs
index 5fc2a6ad2..8e5f02c56 100644
--- a/noodles-sam/src/record/sequence.rs
+++ b/noodles-sam/src/record/sequence.rs
@@ -3,7 +3,7 @@
 pub struct Sequence<'a>(&'a [u8]);
 
 impl<'a> Sequence<'a> {
-    pub(super) fn new(buf: &'a [u8]) -> Self {
+    pub(crate) fn new(buf: &'a [u8]) -> Self {
         Self(buf)
     }
 
@@ -42,6 +42,24 @@ impl<'a> crate::alignment::record::Sequence for Sequence<'a> {
         self.get(i)
     }
 
+    fn split_at_checked(
+        &self,
+        mid: usize,
+    ) -> Option<(
+        Box<dyn crate::alignment::record::Sequence + '_>,
+        Box<dyn crate::alignment::record::Sequence + '_>,
+    )> {
+        if mid <= self.len() {
+            let (left, right) = self.0.split_at(mid);
+            Some((
+                Box::new(Sequence::new(left)),
+                Box::new(Sequence::new(right)),
+            ))
+        } else {
+            None
+        }
+    }
+
     fn iter(&self) -> Box<dyn Iterator<Item = u8> + '_> {
         Box::new(self.as_ref().iter().copied())
     }