use std::{collections::VecDeque, task::Poll};
use futures::Stream;
use futures::future::try_join_all;
use serde::{Deserialize, Serialize};
use turso::Rows;
use crate::RawMessage;
use crate::Result;
use crate::TypedMessage;
use crate::message::filter::MessageFilter;
use crate::storage::Db3File;
pub trait ReadIterator: Sized {
fn from_rows(rows: VecDeque<Rows>) -> Self;
}
pub struct MessageReader<'a> {
connections: &'a [Db3File],
}
impl<'a> MessageReader<'a> {
pub async fn read_all<I: ReadIterator>(&self) -> Result<I> {
self.read_with_filter(&MessageFilter::NONE).await
}
pub async fn read<I: ReadIterator>(&self, message_filter: MessageFilter<'_>) -> Result<I> {
self.read_with_filter(&message_filter).await
}
async fn read_with_filter<I: ReadIterator>(&self, filter: &MessageFilter<'_>) -> Result<I> {
let f = async |db3: &Db3File| db3.query_messages(filter).await;
let rows = self.run_on_connections(f).await?;
Ok(I::from_rows(rows))
}
async fn run_on_connections<F: AsyncFn(&Db3File) -> Result<Rows>>(
&self,
f: F,
) -> Result<VecDeque<Rows>> {
let rows = try_join_all(self.connections.iter().map(|db3| f(db3))).await?;
Ok(rows.into())
}
pub fn create(connections: &'a crate::storage::ConnectionType) -> MessageReader<'a> {
Self {
connections: connections.as_slice(),
}
}
}
pub struct RawReadIterator {
current_rows: Option<Rows>,
results: VecDeque<Rows>,
}
impl ReadIterator for RawReadIterator {
fn from_rows(mut value: VecDeque<Rows>) -> Self {
Self {
current_rows: value.pop_front(),
results: value,
}
}
}
impl Stream for RawReadIterator {
type Item = Result<RawMessage>;
fn poll_next(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if this.current_rows.is_none() {
this.current_rows = this.results.pop_front();
}
let Some(rows) = this.current_rows.as_mut() else {
return Poll::Ready(None);
};
let poll = {
let fut = rows.next();
let mut fut = std::pin::pin!(fut);
fut.as_mut().poll(cx)
};
match poll {
Poll::Ready(Ok(Some(row))) => {
return Poll::Ready(Some(RawMessage::from_row(row)));
}
Poll::Ready(Ok(None)) => {
this.current_rows = None;
}
Poll::Ready(Err(e)) => return Poll::Ready(Some(Err(e.into()))),
Poll::Pending => return Poll::Pending,
}
}
}
}
pub struct TypedReadIterator<T: Serialize + for<'de> Deserialize<'de>> {
current_rows: Option<Rows>,
results: VecDeque<Rows>,
_marker: std::marker::PhantomData<T>,
}
impl<T: Serialize + for<'de> Deserialize<'de>> Unpin for TypedReadIterator<T> {}
impl<T: Serialize + for<'de> Deserialize<'de>> ReadIterator for TypedReadIterator<T> {
fn from_rows(mut value: VecDeque<Rows>) -> Self {
Self {
current_rows: value.pop_front(),
results: value,
_marker: std::marker::PhantomData,
}
}
}
impl<T> Stream for TypedReadIterator<T>
where
T: Serialize + for<'de> Deserialize<'de>,
{
type Item = Result<TypedMessage<T>>;
fn poll_next(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if this.current_rows.is_none() {
this.current_rows = this.results.pop_front();
}
let Some(rows) = this.current_rows.as_mut() else {
return Poll::Ready(None);
};
let poll = {
let fut = rows.next();
let mut fut = std::pin::pin!(fut);
fut.as_mut().poll(cx)
};
match poll {
Poll::Ready(Ok(Some(row))) => {
return Poll::Ready(Some(TypedMessage::<T>::from_row_cdr(row)));
}
Poll::Ready(Ok(None)) => {
this.current_rows = None;
}
Poll::Ready(Err(e)) => return Poll::Ready(Some(Err(e.into()))),
Poll::Pending => return Poll::Pending,
}
}
}
}
#[cfg(test)]
mod test {
use futures::StreamExt;
use super::*;
use crate::Rosbag;
use std::env::current_dir;
#[derive(Debug, Serialize, Deserialize)]
struct Time {
sec: i32,
nanosec: u32,
}
#[derive(Debug, Serialize, Deserialize)]
struct Header {
stamp: Time,
frame_id: String,
}
#[derive(Debug, Serialize, Deserialize)]
struct Image {
header: Header,
height: u32,
width: u32,
encoding: String,
is_bigendian: u8,
step: u32,
data: Vec<u8>,
}
async fn open_test_bag() -> Rosbag {
let mut p = current_dir().expect("Could not get cwd");
p.push("test_bag");
Rosbag::open(p).await.expect("Could not open rosbag")
}
#[tokio::test]
async fn test_iter_typed() {
let rosbag = open_test_bag().await;
let reader = rosbag.read();
let access_request = MessageFilter::new()
.between(100..)
.topics(["/image_raw"])
.expect("Could not construct MessageFilter");
let mut it: TypedReadIterator<Image> =
reader.read(access_request).await.expect("Reading failed");
while let Some(reading_result) = it.next().await {
let msg = reading_result.expect("Reading next message failed");
assert_eq!(msg.data.encoding, "yuv422_yuy2");
println!("{:?}", msg)
}
}
#[tokio::test]
async fn test_iter_raw() {
let rosbag = open_test_bag().await;
let reader = rosbag.read();
let access_request = MessageFilter::new()
.topics(["/image_raw"])
.expect("Could not construct MessageFilter");
let mut it: RawReadIterator = reader.read(access_request).await.expect("Reading failed");
let mut count = 0;
while let Some(reading_result) = it.next().await {
let msg = reading_result.expect("Reading next message failed");
assert_eq!(msg.serialization_format, "cdr");
assert!(!msg.data.is_empty());
count += 1;
}
assert_eq!(count, 214, "Did not read expected number of messages");
}
}