Importer V3 (#395)

* start of importer refactor

* fish

* resh

* zsh

1 files changed, 85 insertions, 14 deletions

diff --git a/atuin-client/src/import/mod.rs b/atuin-client/src/import/mod.rs
index 8d4aa17f..07178d17 100644
--- a/atuin-client/src/import/mod.rs
+++ b/atuin-client/src/import/mod.rs
@@ -1,9 +1,8 @@
-use std::{
-    io::{BufRead, BufReader, Read, Seek, SeekFrom},
-    path::{Path, PathBuf},
-};
+use std::path::PathBuf;
 
-use eyre::Result;
+use async_trait::async_trait;
+use eyre::{bail, Result};
+use memchr::Memchr;
 
 use crate::history::History;
 
@@ -12,16 +11,88 @@ pub mod fish;
 pub mod resh;
 pub mod zsh;
 
-// this could probably be sped up
-fn count_lines(buf: &mut BufReader<impl Read + Seek>) -> Result<usize> {
-    let lines = buf.lines().count();
-    buf.seek(SeekFrom::Start(0))?;
+#[async_trait]
+pub trait Importer: Sized {
+    const NAME: &'static str;
+    async fn new() -> Result<Self>;
+    async fn entries(&mut self) -> Result<usize>;
+    async fn load(self, loader: &mut impl Loader) -> Result<()>;
+}
 
-    Ok(lines)
+#[async_trait]
+pub trait Loader: Sync + Send {
+    async fn push(&mut self, hist: History) -> eyre::Result<()>;
 }
 
-pub trait Importer: IntoIterator<Item = Result<History>> + Sized {
-    const NAME: &'static str;
-    fn histpath() -> Result<PathBuf>;
-    fn parse(path: impl AsRef<Path>) -> Result<Self>;
+fn unix_byte_lines(input: &[u8]) -> impl Iterator<Item = &[u8]> {
+    UnixByteLines {
+        iter: memchr::memchr_iter(b'\n', input),
+        bytes: input,
+        i: 0,
+    }
+}
+
+struct UnixByteLines<'a> {
+    iter: Memchr<'a>,
+    bytes: &'a [u8],
+    i: usize,
+}
+
+impl<'a> Iterator for UnixByteLines<'a> {
+    type Item = &'a [u8];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let j = self.iter.next()?;
+        let out = &self.bytes[self.i..j];
+        self.i = j + 1;
+        Some(out)
+    }
+
+    fn count(self) -> usize
+    where
+        Self: Sized,
+    {
+        self.iter.count()
+    }
+}
+
+fn count_lines(input: &[u8]) -> usize {
+    unix_byte_lines(input).count()
+}
+
+fn get_histpath<D>(def: D) -> Result<PathBuf>
+where
+    D: FnOnce() -> Result<PathBuf>,
+{
+    if let Ok(p) = std::env::var("HISTFILE") {
+        is_file(PathBuf::from(p))
+    } else {
+        is_file(def()?)
+    }
+}
+
+fn is_file(p: PathBuf) -> Result<PathBuf> {
+    if p.is_file() {
+        Ok(p)
+    } else {
+        bail!("Could not find history file {:?}. Try setting $HISTFILE", p)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[derive(Default)]
+    pub struct TestLoader {
+        pub buf: Vec<History>,
+    }
+
+    #[async_trait]
+    impl Loader for TestLoader {
+        async fn push(&mut self, hist: History) -> Result<()> {
+            self.buf.push(hist);
+            Ok(())
+        }
+    }
 }