Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions src/derive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,18 @@ pub struct DerivedKey {
pub p2pkh_uncompressed: String,
/// P2WPKH (bech32) address
pub p2wpkh: String,
/// P2TR (Taproot bech32m) address
pub p2tr: String,
}

impl DerivedKey {
/// Get all addresses as slice for matching.
pub fn addresses(&self) -> [&str; 3] {
pub fn addresses(&self) -> [&str; 4] {
[
&self.p2pkh_compressed,
&self.p2pkh_uncompressed,
&self.p2wpkh,
&self.p2tr,
]
}
}
Expand Down Expand Up @@ -117,6 +120,10 @@ impl KeyDeriver {
.expect("valid compressed pubkey");
let p2wpkh = Address::p2wpkh(&compressed_pk, self.network).to_string();

// P2TR (Taproot, key-path spend with no script tree)
let (x_only_key, _parity) = secp_pubkey.x_only_public_key();
let p2tr = Address::p2tr(&self.secp, x_only_key, None, self.network).to_string();

// Decimal representation (big-endian)
let private_key_decimal = BigUint::from_bytes_be(&key_bytes).to_string();

Expand Down Expand Up @@ -161,6 +168,7 @@ impl KeyDeriver {
p2pkh_compressed,
p2pkh_uncompressed,
p2wpkh,
p2tr,
}
}
}
Expand Down Expand Up @@ -206,10 +214,11 @@ mod tests {
let derived = deriver.derive(&key);

let addrs = derived.addresses();
assert_eq!(addrs.len(), 3);
assert_eq!(addrs.len(), 4);
assert!(addrs[0].starts_with('1')); // P2PKH compressed
assert!(addrs[1].starts_with('1')); // P2PKH uncompressed
assert!(addrs[2].starts_with("bc1q")); // P2WPKH
assert!(addrs[3].starts_with("bc1p")); // P2TR
}

#[test]
Expand Down
26 changes: 26 additions & 0 deletions src/matcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ pub enum AddressType {
P2pkhCompressed,
P2pkhUncompressed,
P2wpkh,
P2tr,
}

impl AddressType {
Expand All @@ -31,6 +32,7 @@ impl AddressType {
AddressType::P2pkhCompressed => "p2pkh_compressed",
AddressType::P2pkhUncompressed => "p2pkh_uncompressed",
AddressType::P2wpkh => "p2wpkh",
AddressType::P2tr => "p2tr",
}
}
}
Expand Down Expand Up @@ -94,6 +96,14 @@ impl Matcher {
});
}

// Check P2TR (Taproot)
if self.targets.contains(&derived.p2tr) {
return Some(MatchInfo {
address_type: AddressType::P2tr,
address: derived.p2tr.clone(),
});
}

None
}

Expand Down Expand Up @@ -137,6 +147,22 @@ mod tests {
assert_eq!(info.address, "1JwSSubhmg6iPtRjtyqhUYYH7bZg3Lfy1T");
}

#[test]
fn test_matcher_p2tr() {
let key = [1u8; 32];
let deriver = KeyDeriver::new();
let derived = deriver.derive(&key);

let matcher = Matcher::from_addresses(vec![derived.p2tr.clone()]);

let result = matcher.check(&derived);
assert!(result.is_some());

let info = result.unwrap();
assert!(matches!(info.address_type, AddressType::P2tr));
assert!(info.address.starts_with("bc1p"));
}

#[test]
fn test_matcher_no_match() {
let key = [1u8; 32];
Expand Down
3 changes: 3 additions & 0 deletions src/output/console.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ impl Output for ConsoleOutput {
writeln!(w, "p2pkh_compressed: {}", derived.p2pkh_compressed)?;
writeln!(w, "p2pkh_uncompressed: {}", derived.p2pkh_uncompressed)?;
writeln!(w, "p2wpkh: {}", derived.p2wpkh)?;
writeln!(w, "p2tr: {}", derived.p2tr)?;
} else {
// Compact format: source,transform,privkey,address_compressed
writeln!(
Expand Down Expand Up @@ -115,6 +116,7 @@ impl Output for ConsoleOutput {
writeln!(w, "P2PKH (compressed): {}", derived.p2pkh_compressed)?;
writeln!(w, "P2PKH (uncompressed): {}", derived.p2pkh_uncompressed)?;
writeln!(w, "P2WPKH: {}", derived.p2wpkh)?;
writeln!(w, "P2TR: {}", derived.p2tr)?;
writeln!(w, "=========================")?;

Ok(())
Expand Down Expand Up @@ -165,6 +167,7 @@ mod tests {
p2pkh_compressed: "1Address".to_string(),
p2pkh_uncompressed: "1Uncompressed".to_string(),
p2wpkh: "bc1q...".to_string(),
p2tr: "bc1p...".to_string(),
}
}

Expand Down
1 change: 1 addition & 0 deletions src/output/multi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ mod tests {
p2pkh_compressed: "1ABC".to_string(),
p2pkh_uncompressed: "1DEF".to_string(),
p2wpkh: "bc1q".to_string(),
p2tr: "bc1p".to_string(),
}
}

Expand Down
5 changes: 5 additions & 0 deletions src/output/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,10 @@ impl StorageOutput {
address_type: "p2wpkh",
address: &derived.p2wpkh,
},
AddressRecord {
address_type: "p2tr",
address: &derived.p2tr,
},
];

let export_formats = [
Expand Down Expand Up @@ -276,6 +280,7 @@ mod tests {
p2pkh_compressed: "1ABC123".to_string(),
p2pkh_uncompressed: "1DEF456".to_string(),
p2wpkh: "bc1qtest".to_string(),
p2tr: "bc1ptest".to_string(),
}
}

Expand Down
8 changes: 7 additions & 1 deletion src/storage/iceberg/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,12 @@ pub fn build_iceberg_schema() -> Result<IcebergSchema> {
Type::Primitive(PrimitiveType::String),
)
.into(),
NestedField::optional(
next_id(),
fields::ADDRESS_P2TR,
Type::Primitive(PrimitiveType::String),
)
.into(),
Comment thread
coderabbitai[bot] marked this conversation as resolved.
NestedField::optional(
next_id(),
fields::WIF_COMPRESSED,
Expand Down Expand Up @@ -150,7 +156,7 @@ mod tests {
#[test]
fn build_schema_succeeds() {
let schema = build_iceberg_schema().unwrap();
assert_eq!(schema.as_struct().fields().len(), 19);
assert_eq!(schema.as_struct().fields().len(), 20);
}

#[test]
Expand Down
4 changes: 2 additions & 2 deletions src/storage/parquet_backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -337,7 +337,7 @@ mod tests {
fn schema_returns_result_schema() {
let backend = ParquetBackend::new("results", "sha256");
let schema = backend.schema();
assert_eq!(schema.fields().len(), 19);
assert_eq!(schema.fields().len(), 20);
assert_eq!(schema.field(0).name(), "source");
}

Expand Down Expand Up @@ -581,7 +581,7 @@ mod tests {

let batch = &batches[0];
assert_eq!(batch.num_rows(), 2);
assert_eq!(batch.num_columns(), 19);
assert_eq!(batch.num_columns(), 20);

let source_col = batch
.column(0)
Expand Down
29 changes: 20 additions & 9 deletions src/storage/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ pub mod fields {
pub const ADDRESS_P2PKH_COMPRESSED: &str = "address_p2pkh_compressed";
pub const ADDRESS_P2PKH_UNCOMPRESSED: &str = "address_p2pkh_uncompressed";
pub const ADDRESS_P2WPKH: &str = "address_p2wpkh";
pub const ADDRESS_P2TR: &str = "address_p2tr";

// Export formats
pub const WIF_COMPRESSED: &str = "wif_compressed";
Expand All @@ -52,7 +53,7 @@ pub mod fields {
/// Variable-length fields (public_keys, addresses, export_formats) are mapped
/// to fixed columns based on known Bitcoin formats.
///
/// # Schema (19 columns)
/// # Schema (20 columns)
///
/// | Column | Type | Nullable | Description |
/// |--------|------|----------|-------------|
Expand All @@ -73,6 +74,7 @@ pub mod fields {
/// | address_p2pkh_compressed | Utf8 | Yes | P2PKH (compressed) |
/// | address_p2pkh_uncompressed | Utf8 | Yes | P2PKH (uncompressed) |
/// | address_p2wpkh | Utf8 | Yes | P2WPKH (native segwit) |
/// | address_p2tr | Utf8 | Yes | P2TR (Taproot) |
/// | wif_compressed | Utf8 | Yes | WIF compressed |
/// | wif_uncompressed | Utf8 | Yes | WIF uncompressed |
pub fn result_schema() -> Schema {
Expand Down Expand Up @@ -106,6 +108,7 @@ pub fn result_schema() -> Schema {
Field::new(fields::ADDRESS_P2PKH_COMPRESSED, DataType::Utf8, true),
Field::new(fields::ADDRESS_P2PKH_UNCOMPRESSED, DataType::Utf8, true),
Field::new(fields::ADDRESS_P2WPKH, DataType::Utf8, true),
Field::new(fields::ADDRESS_P2TR, DataType::Utf8, true),
// Export formats (nullable)
Field::new(fields::WIF_COMPRESSED, DataType::Utf8, true),
Field::new(fields::WIF_UNCOMPRESSED, DataType::Utf8, true),
Expand Down Expand Up @@ -232,6 +235,12 @@ pub fn records_to_batch(records: &[ResultRecord<'_>]) -> Result<RecordBatch, Arr
.map(|r| find_address(r.addresses, "p2wpkh"))
.collect::<Vec<_>>(),
));
let address_p2tr_array: ArrayRef = Arc::new(StringArray::from(
records
.iter()
.map(|r| find_address(r.addresses, "p2tr"))
.collect::<Vec<_>>(),
));

// Export format arrays (nullable)
let wif_compressed_array: ArrayRef = Arc::new(StringArray::from(
Expand Down Expand Up @@ -267,6 +276,7 @@ pub fn records_to_batch(records: &[ResultRecord<'_>]) -> Result<RecordBatch, Arr
address_p2pkh_compressed_array,
address_p2pkh_uncompressed_array,
address_p2wpkh_array,
address_p2tr_array,
wif_compressed_array,
wif_uncompressed_array,
],
Expand All @@ -281,9 +291,9 @@ mod tests {
use arrow::datatypes::DataType;

#[test]
fn schema_has_19_fields() {
fn schema_has_20_fields() {
let schema = result_schema();
assert_eq!(schema.fields().len(), 19);
assert_eq!(schema.fields().len(), 20);
}

#[test]
Expand Down Expand Up @@ -311,6 +321,7 @@ mod tests {
"address_p2pkh_compressed",
"address_p2pkh_uncompressed",
"address_p2wpkh",
"address_p2tr",
"wif_compressed",
"wif_uncompressed",
]
Expand All @@ -337,7 +348,7 @@ mod tests {
assert_eq!(schema.field(10).data_type(), &DataType::UInt16);
assert_eq!(schema.field(11).data_type(), &DataType::UInt8);

for i in 12..19 {
for i in 12..20 {
assert_eq!(schema.field(i).data_type(), &DataType::Utf8);
}
}
Expand All @@ -355,7 +366,7 @@ mod tests {
);
}

let nullable = [4, 12, 13, 14, 15, 16, 17, 18];
let nullable = [4, 12, 13, 14, 15, 16, 17, 18, 19];
for i in nullable {
assert!(
schema.field(i).is_nullable(),
Expand All @@ -369,7 +380,7 @@ mod tests {
fn records_to_batch_empty() {
let batch = records_to_batch(&[]).unwrap();
assert_eq!(batch.num_rows(), 0);
assert_eq!(batch.num_columns(), 19);
assert_eq!(batch.num_columns(), 20);
assert_eq!(batch.schema(), Arc::new(result_schema()));
}

Expand Down Expand Up @@ -449,7 +460,7 @@ mod tests {
let batch = records_to_batch(&[record]).unwrap();

assert_eq!(batch.num_rows(), 1);
assert_eq!(batch.num_columns(), 19);
assert_eq!(batch.num_columns(), 20);

let source_col = batch
.column(0)
Expand Down Expand Up @@ -501,7 +512,7 @@ mod tests {
assert_eq!(addr_p2wpkh_col.value(0), "bc1qtest");

let wif_compressed_col = batch
.column(17)
.column(18)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
Expand Down Expand Up @@ -558,7 +569,7 @@ mod tests {
assert!(addr_p2pkh.is_null(0));

let wif = batch
.column(17)
.column(18)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
Expand Down
Loading