feat: add name and namespace parsing in stx-genesis

stx-genesis/build.rs

@@ -22,50 +22,38 @@ fn open_chainstate_file() -> File {
     File::open("chainstate.txt").unwrap()
 }
 
-pub fn write_archives() -> std::io::Result<()> {
+fn write_archive(output_file_name: &str, section_name: &str) -> std::io::Result<()> {
     let out_dir = env::var_os("OUT_DIR").unwrap();
+    let chainstate_file = open_chainstate_file();
+    let reader = BufReader::new(chainstate_file);
+    let out_file_path = Path::new(&out_dir).join(output_file_name);
+    let out_file = File::create(out_file_path)?;
+    let mut encoder = deflate::Encoder::new(out_file);
+
+    let section_header = "-----BEGIN ".to_owned() + section_name + "-----";
+    let section_footer = "-----END ".to_owned() + section_name + "-----";
+
+    for line in reader
+        .lines()
+        .map(|line| line.unwrap())
+        .skip_while(|line| !line.eq(&section_header))
+        // skip table header line
+        .skip(2)
+        .take_while(|line| !line.eq(&section_footer))
     {
-        let chainstate_file = open_chainstate_file();
-        let reader = BufReader::new(chainstate_file);
-        let balances_file_path = Path::new(&out_dir).join("account_balances.gz");
-        let balances_file = File::create(balances_file_path)?;
-        let mut balances_encoder = deflate::Encoder::new(balances_file);
-
-        for line in reader
-            .lines()
-            .map(|line| line.unwrap())
-            .skip_while(|line| !line.eq(&"-----BEGIN STX BALANCES-----"))
-            // skip table header line "address,balance"
-            .skip(2)
-            .take_while(|line| !line.eq(&"-----END STX BALANCES-----"))
-        {
-            balances_encoder.write_all(&[line.as_bytes(), &[b'\n']].concat())?;
-        }
-
-        let mut balances_file = balances_encoder.finish().into_result().unwrap();
-        balances_file.flush()?;
+        encoder.write_all(&[line.as_bytes(), &[b'\n']].concat())?;
     }
-    {
-        let chainstate_file = open_chainstate_file();
-        let reader = BufReader::new(chainstate_file);
-        let vesting_file_path = Path::new(&out_dir).join("account_vesting.gz");
-        let vesting_file = File::create(vesting_file_path)?;
-        let mut vesting_encoder = deflate::Encoder::new(vesting_file);
 
-        for line in reader
-            .lines()
-            .map(|line| line.unwrap())
-            .skip_while(|line| !line.eq(&"-----BEGIN STX VESTING-----"))
-            // skip table header line "address,value,blocks"
-            .skip(2)
-            .take_while(|line| !line.eq(&"-----END STX VESTING-----"))
-        {
-            vesting_encoder.write_all(&[line.as_bytes(), &[b'\n']].concat())?;
-        }
+    let mut out_file = encoder.finish().into_result().unwrap();
+    out_file.flush()?;
+    Ok(())
+}
 
-        let mut vesting_file = vesting_encoder.finish().into_result().unwrap();
-        vesting_file.flush()?;
-    }
+pub fn write_archives() -> std::io::Result<()> {
+    write_archive("account_balances.gz", "STX BALANCES")?;
+    write_archive("account_vesting.gz", "STX VESTING")?;
+    write_archive("namespaces.gz", "NAMESPACES")?;
+    write_archive("names.gz", "NAMES")?;
     Ok(())
 }
 

stx-genesis/src/lib.rs

@@ -19,41 +19,120 @@ pub struct GenesisAccountVesting {
     pub block_height: u64,
 }
 
+pub struct GenesisNamespace {
+    pub namespace_id: String,
+    pub address: String,
+    pub reveal_block: i64,
+    pub ready_block: i64,
+    pub buckets: String,
+    pub base: String,
+    pub coeff: String,
+    pub nonalpha_discount: String,
+    pub no_vowel_discount: String,
+    pub lifetime: String,
+}
+
+pub struct GenesisName {
+    pub name: String,
+    pub address: String,
+    pub registered_at: i64,
+    pub expire_block: i64,
+    pub zonefile_hash: String,
+}
+
 pub static GENESIS_CHAINSTATE_HASH: &str =
     include_str!(concat!(env!("OUT_DIR"), "/chainstate.txt.sha256"));
 
+fn iter_deflated_csv(deflate_bytes: &'static [u8]) -> Box<dyn Iterator<Item = Vec<String>>> {
+    let cursor = io::Cursor::new(deflate_bytes);
+    let deflate_decoder = deflate::Decoder::new(cursor);
+    let buff_reader = BufReader::new(deflate_decoder);
+    let line_iter = buff_reader
+        .lines()
+        .map(|line| line.unwrap())
+        .map(|line| line.split(",").map(String::from).collect());
+    return Box::new(line_iter);
+}
+
 pub fn read_balances() -> Box<dyn Iterator<Item = GenesisAccountBalance>> {
     let account_balances_bytes = include_bytes!(concat!(env!("OUT_DIR"), "/account_balances.gz"));
-    let cursor = io::Cursor::new(account_balances_bytes);
-    let balances_encoder = deflate::Decoder::new(cursor);
-    let buff_reader = BufReader::new(balances_encoder);
-    let balances = buff_reader.lines().map(|line| line.unwrap()).map(|line| {
-        let mut parts = line.split(",");
-        let addr = parts.next().unwrap();
-        let balance = parts.next().unwrap().parse::<u64>().unwrap();
-        GenesisAccountBalance {
-            address: addr.to_string(),
-            amount: balance,
-        }
+    let balances = iter_deflated_csv(account_balances_bytes).map(|cols| GenesisAccountBalance {
+        address: cols[0].to_string(),
+        amount: cols[1].parse::<u64>().unwrap(),
     });
     return Box::new(balances);
 }
 
 pub fn read_vesting() -> Box<dyn Iterator<Item = GenesisAccountVesting>> {
     let account_balances_bytes = include_bytes!(concat!(env!("OUT_DIR"), "/account_vesting.gz"));
-    let cursor = io::Cursor::new(account_balances_bytes);
-    let balances_encoder = deflate::Decoder::new(cursor);
-    let buff_reader = BufReader::new(balances_encoder);
-    let balances = buff_reader.lines().map(|line| line.unwrap()).map(|line| {
-        let mut parts = line.split(",");
-        let addr = parts.next().unwrap();
-        let amount = parts.next().unwrap().parse::<u64>().unwrap();
-        let block_height = parts.next().unwrap().parse::<u64>().unwrap();
-        GenesisAccountVesting {
-            address: addr.to_string(),
-            amount: amount,
-            block_height: block_height,
-        }
+    let vesting = iter_deflated_csv(account_balances_bytes).map(|cols| GenesisAccountVesting {
+        address: cols[0].to_string(),
+        amount: cols[1].parse::<u64>().unwrap(),
+        block_height: cols[2].parse::<u64>().unwrap(),
     });
-    return Box::new(balances);
+    return Box::new(vesting);
+}
+
+pub fn read_namespaces() -> Box<dyn Iterator<Item = GenesisNamespace>> {
+    let namespaces_bytes = include_bytes!(concat!(env!("OUT_DIR"), "/namespaces.gz"));
+    let namespaces = iter_deflated_csv(namespaces_bytes).map(|cols| GenesisNamespace {
+        namespace_id: cols[0].to_string(),
+        address: cols[1].to_string(),
+        reveal_block: cols[2].parse::<i64>().unwrap(),
+        ready_block: cols[3].parse::<i64>().unwrap(),
+        buckets: cols[4].to_string(),
+        base: cols[5].to_string(),
+        coeff: cols[6].to_string(),
+        nonalpha_discount: cols[7].to_string(),
+        no_vowel_discount: cols[8].to_string(),
+        lifetime: cols[9].to_string(),
+    });
+    return Box::new(namespaces);
+}
+
+pub fn read_names() -> Box<dyn Iterator<Item = GenesisName>> {
+    let names_bytes = include_bytes!(concat!(env!("OUT_DIR"), "/names.gz"));
+    let names = iter_deflated_csv(names_bytes).map(|cols| GenesisName {
+        name: cols[0].to_string(),
+        address: cols[1].to_string(),
+        registered_at: cols[2].parse::<i64>().unwrap(),
+        expire_block: cols[3].parse::<i64>().unwrap(),
+        zonefile_hash: cols[4].to_string(),
+    });
+    return Box::new(names);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Test the decompression and line parsing
+
+    #[test]
+    fn test_balances_read() {
+        for balance in read_balances() {
+            assert!(balance.amount > 0);
+        }
+    }
+
+    #[test]
+    fn test_vestings_read() {
+        for vesting in read_vesting() {
+            assert!(vesting.amount > 0);
+        }
+    }
+
+    #[test]
+    fn test_namespaces_read() {
+        for namespace in read_namespaces() {
+            assert!(namespace.ready_block > 0);
+        }
+    }
+
+    #[test]
+    fn test_names_read() {
+        for name in read_names() {
+            assert!(name.registered_at > 0);
+        }
+    }
 }