implement aft-head

3 files changed, 242 insertions, 5 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 80a6d46..e998284 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8,6 +8,8 @@ dependencies = [
  "assert_cmd",
  "atty",
  "colored",
+ "serde_json",
+ "structopt",
  "tabwriter",
 ]
 
@@ -42,6 +44,23 @@ dependencies = [
 ]
 
 [[package]]
+name = "bitflags"
+version = "1.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf1de2fe8c75bc145a2f577add951f8134889b4795d47466a54a5c846d691693"
+
+[[package]]
+name = "clap"
+version = "2.33.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bdfa80d47f954d53a35a64987ca1422f495b8d6483c0fe9f7117b36c2a792129"
+dependencies = [
+ "bitflags",
+ "textwrap",
+ "unicode-width",
+]
+
+[[package]]
 name = "colored"
 version = "1.9.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -65,6 +84,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10"
 
 [[package]]
+name = "heck"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20564e78d53d2bb135c343b3f47714a56af2061f1c928fdb541dc7b9fdd94205"
+dependencies = [
+ "unicode-segmentation",
+]
+
+[[package]]
 name = "hermit-abi"
 version = "0.1.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -74,6 +102,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "itoa"
+version = "0.4.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8b7a7c0c47db5545ed3fef7468ee7bb5b74691498139e4b3f6a20685dc6dd8e"
+
+[[package]]
 name = "lazy_static"
 version = "1.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -112,6 +146,50 @@ dependencies = [
 ]
 
 [[package]]
+name = "proc-macro-error"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "98e9e4b82e0ef281812565ea4751049f1bdcdfccda7d3f459f2e138a40c08678"
+dependencies = [
+ "proc-macro-error-attr",
+ "proc-macro2",
+ "quote",
+ "syn",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro-error-attr"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4f5444ead4e9935abd7f27dc51f7e852a0569ac888096d5ec2499470794e2e53"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+ "syn-mid",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "beae6331a816b1f65d04c45b078fd8e6c93e8071771f41b8163255bbd8d7c8fa"
+dependencies = [
+ "unicode-xid",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54a21852a652ad6f610c9510194f398ff6f8692e334fd1145fed931f7fbe44ea"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
 name = "regex"
 version = "1.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -127,6 +205,75 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7fe5bd57d1d7414c6b5ed48563a2c855d995ff777729dcd91c369ec7fea395ae"
 
 [[package]]
+name = "ryu"
+version = "1.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "71d301d4193d031abdd79ff7e3dd721168a9572ef3fe51a1517aba235bd8f86e"
+
+[[package]]
+name = "serde"
+version = "1.0.111"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c9124df5b40cbd380080b2cc6ab894c040a3070d995f5c9dc77e18c34a8ae37d"
+
+[[package]]
+name = "serde_json"
+version = "1.0.53"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "993948e75b189211a9b31a7528f950c6adc21f9720b6438ff80a7fa2f864cea2"
+dependencies = [
+ "itoa",
+ "ryu",
+ "serde",
+]
+
+[[package]]
+name = "structopt"
+version = "0.3.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "863246aaf5ddd0d6928dfeb1a9ca65f505599e4e1b399935ef7e75107516b4ef"
+dependencies = [
+ "clap",
+ "lazy_static",
+ "structopt-derive",
+]
+
+[[package]]
+name = "structopt-derive"
+version = "0.4.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d239ca4b13aee7a2142e6795cbd69e457665ff8037aed33b3effdc430d2f927a"
+dependencies = [
+ "heck",
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "syn"
+version = "1.0.30"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93a56fabc59dce20fe48b6c832cc249c713e7ed88fa28b0ee0a3bfcaae5fe4e2"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-xid",
+]
+
+[[package]]
+name = "syn-mid"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7be3539f6c128a931cf19dcee741c1af532c7fd387baa739c03dd2e96479338a"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
 name = "tabwriter"
 version = "1.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -138,18 +285,45 @@ dependencies = [
 ]
 
 [[package]]
+name = "textwrap"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d326610f408c7a4eb6f51c37c330e496b08506c9457c9d34287ecc38809fb060"
+dependencies = [
+ "unicode-width",
+]
+
+[[package]]
 name = "treeline"
 version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a7f741b240f1a48843f9b8e0444fb55fb2a4ff67293b50a9179dfd5ea67f8d41"
 
 [[package]]
+name = "unicode-segmentation"
+version = "1.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e83e153d1053cbb5a118eeff7fd5be06ed99153f00dbcd8ae310c5fb2b22edc0"
+
+[[package]]
 name = "unicode-width"
 version = "0.1.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "caaa9d531767d1ff2150b9332433f32a24622147e5ebb1f26409d5da67afd479"
 
 [[package]]
+name = "unicode-xid"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "826e7639553986605ec5979c7dd957c7895e93eabed50ab2ffa7f6128a75097c"
+
+[[package]]
+name = "version_check"
+version = "0.9.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b5a972e5669d67ba988ce3dc826706fb0a8b01471c088cb0b6110b805cc36aed"
+
+[[package]]
 name = "wait-timeout"
 version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index e5428b2..e2937b8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,9 +5,9 @@ authors = ["Bryan Newbold <bnewbold@robocracy.org>"]
 edition = "2018"
 
 [dependencies]
-#serde_json = "1.0"
 #bstr = "0.2"
-#structopt = "0.3"
+serde_json = "1.0"
+structopt = { version = "0.3", default-features = false }
 anyhow = "1.0"
 colored = "1.9"
 atty = "0.2"
diff --git a/src/bin/aft-head.rs b/src/bin/aft-head.rs
index b8c848c..3fecfa2 100644
--- a/src/bin/aft-head.rs
+++ b/src/bin/aft-head.rs
@@ -1,8 +1,71 @@
-
 /*
  * Takes the first n lines of input, passes to stdout.
  */
 
-fn main() {
-    println!("Hello, world2!");
+use std::error::Error;
+use std::io::{self, Read, Write, BufRead};
+use structopt::StructOpt;
+use std::path::PathBuf;
+
+#[derive(StructOpt)]
+#[structopt(rename_all = "kebab-case", about = "take first N rows of input (AFT-aware)" )]
+struct Opt {
+    #[structopt(long = "--lines", short = "-n", default_value = "10")]
+    lines: u64,
+
+    #[structopt(parse(from_os_str))]
+    input: Option<PathBuf>,
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+
+    let opt = Opt::from_args();
+
+    // treat "-" as "use stdin"
+    let input = match opt.input {
+        Some(s) if s.to_string_lossy() == "-" => None,
+        other => other,
+    };
+    let stdin = io::stdin();
+    let stdin = stdin.lock();
+    let mut input: Box<dyn BufRead> = match input {
+        None => {
+            Box::new(io::BufReader::new(stdin))
+        },
+        Some(path) => {
+            let f = std::fs::File::open(path)?;
+            Box::new(io::BufReader::new(f))
+        },
+    };
+    let mut stdout = io::stdout();
+
+    // read first byte of input to check if we are getting AFT; if not just pass through
+    let mut first_byte = [0; 1];
+    let got = input.read(&mut first_byte)?;
+    if got != 1 {
+        panic!("couldn't read a byte from input");
+    };
+
+    let is_aft = first_byte[0] == 0x01;
+    if !is_aft {
+        // simple implementation of lines
+        stdout.write(&first_byte)?;
+        for l in input.lines().take(opt.lines as usize) {
+            writeln!(stdout, "{}", l?)?;
+        }
+        return Ok(())
+    };
+
+    let mut header_bytes = vec![];
+    // TODO: sanity limit on length
+    input.read_until(0x02, &mut header_bytes)?;
+    if header_bytes[header_bytes.len()-1] != 0x02 {
+        panic!("expected an AFT header");
+    };
+    stdout.write(&first_byte)?;
+    stdout.write(&header_bytes)?;
+    for l in input.lines().take(opt.lines as usize) {
+        writeln!(stdout, "{}", l?)?;
+    }
+    Ok(())
 }