1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::{ids::short, key, secp256k1fx};

/// Support multiple keys as a chain.
/// ref. https://pkg.go.dev/github.com/ava-labs/avalanchego/vms/secp256k1fx#Keychain
/// ref. https://github.com/ava-labs/avalanchego/blob/v1.7.9/wallet/chain/p/builder.go
#[derive(Debug, Serialize, Deserialize, Eq, PartialEq, Clone)]
pub struct Keychain<T: key::ReadOnly + key::SignOnly> {
    pub keys: Vec<T>,
    pub short_addr_to_key_index: HashMap<short::Id, u32>,
}

impl<T> Keychain<T>
where
    T: key::ReadOnly + key::SignOnly + Clone,
{
    pub fn new(keys: Vec<T>) -> Self {
        let mut short_addr_to_key_index = HashMap::new();
        for (pos, k) in keys.iter().enumerate() {
            short_addr_to_key_index.insert(k.get_short_address(), pos as u32);
        }
        Self {
            keys,
            short_addr_to_key_index,
        }
    }

    /// ref. https://pkg.go.dev/github.com/ava-labs/avalanchego/vms/secp256k1fx#Keychain.Get
    pub fn get(&self, short_addr: &short::Id) -> Option<T> {
        self.short_addr_to_key_index
            .get(short_addr)
            .map(|k| self.keys[(*k) as usize].clone())
    }

    /// ref. https://pkg.go.dev/github.com/ava-labs/avalanchego/vms/secp256k1fx#Keychain.Match
    pub fn match_threshold(
        &self,
        output_owners: &secp256k1fx::OutputOwners,
        time: u64,
    ) -> Option<(Vec<u32>, Vec<T>)> {
        if output_owners.locktime > time {
            // output owners are still locked
            return None;
        }

        let mut sig_indices: Vec<u32> = Vec::new();
        let mut keys: Vec<T> = Vec::new();
        for (pos, addr) in output_owners.addrs.iter().enumerate() {
            let key = self.get(addr);
            if key.is_none() {
                continue;
            }
            sig_indices.push(pos as u32);
            keys.push(key.unwrap());

            if (keys.len() as u32) == output_owners.threshold {
                break;
            }
        }

        let n = keys.len();
        if (n as u32) == output_owners.threshold {
            Some((sig_indices, keys))
        } else {
            None
        }
    }

    /// Returns "None" if the threshold is NOT met.
    /// ref. https://pkg.go.dev/github.com/ava-labs/avalanchego/vms/secp256k1fx#Keychain.Spend
    /// TODO: support spend on "secp256k1fx::MintOutput"
    pub fn spend(
        &self,
        output: &secp256k1fx::TransferOutput,
        time: u64,
    ) -> Option<(secp256k1fx::TransferInput, Vec<T>)> {
        let res = self.match_threshold(&output.output_owners, time);
        let threshold_met = res.is_some();
        if !threshold_met {
            return None;
        }

        let (sig_indices, keys) = res.unwrap();
        Some((
            secp256k1fx::TransferInput {
                amount: output.amount,
                sig_indices,
            },
            keys,
        ))
    }
}