Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

optimize wallet tool #11154

Merged
merged 1 commit into from
Apr 13, 2022
Merged
Changes from all commits
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
23 changes: 11 additions & 12 deletions tests/wallet_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
class WalletTool:
next_address = 0
pubkey_num_lookup: Dict[bytes, uint32] = {}
puzzle_pk_cache: Dict[bytes32, PrivateKey] = {}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
puzzle_pk_cache: Dict[bytes32, PrivateKey] = {}
puzzle_sk_cache: Dict[bytes32, PrivateKey] = {}

By cryptography programmer conventions, this variable should be named puzzle_sk_cache

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted to keep my changes focused on just this issue of mine. Also, by the same standards, the PrivateKey type should be called SecretKey :)


def __init__(self, constants: ConsensusConstants, sk: Optional[PrivateKey] = None):
self.constants = constants
Expand All @@ -48,29 +49,27 @@ def get_next_address_index(self) -> uint32:
return self.next_address

def get_private_key_for_puzzle_hash(self, puzzle_hash: bytes32) -> PrivateKey:
if puzzle_hash in self.puzzle_pk_cache:
child = self.puzzle_pk_cache[puzzle_hash]
private = master_sk_to_wallet_sk(self.private_key, uint32(child))
# pubkey = private.get_g1()
return private
else:
for child in range(self.next_address):
pubkey = master_sk_to_wallet_sk(self.private_key, uint32(child)).get_g1()
if puzzle_hash == puzzle_for_pk(bytes(pubkey)).get_tree_hash():
return master_sk_to_wallet_sk(self.private_key, uint32(child))
sk = self.puzzle_pk_cache.get(puzzle_hash)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
sk = self.puzzle_pk_cache.get(puzzle_hash)
sk = self.puzzle_sk_cache.get(puzzle_hash)

if sk:
return sk
for child in range(self.next_address):
pubkey = master_sk_to_wallet_sk(self.private_key, uint32(child)).get_g1()
if puzzle_hash == puzzle_for_pk(bytes(pubkey)).get_tree_hash():
return master_sk_to_wallet_sk(self.private_key, uint32(child))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should you add this one to the cache, too? I don't see a call path where it would add it in this case, in the diff.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think that would help, and maybe make things worse. I was considering removing entries from the cache here. Typically we call this function to get the private key to spend a coin. once spent, we are very unlikely to call this for the same puzzle hash again

raise ValueError(f"Do not have the keys for puzzle hash {puzzle_hash}")

def puzzle_for_pk(self, pubkey: bytes) -> Program:
return puzzle_for_pk(pubkey)

def get_new_puzzle(self) -> Program:
next_address_index: uint32 = self.get_next_address_index()
pubkey: G1Element = master_sk_to_wallet_sk(self.private_key, next_address_index).get_g1()
sk: PrivateKey = master_sk_to_wallet_sk(self.private_key, next_address_index)
pubkey: G1Element = sk.get_g1()
self.pubkey_num_lookup[bytes(pubkey)] = next_address_index

puzzle: Program = puzzle_for_pk(pubkey)

self.puzzle_pk_cache[puzzle.get_tree_hash()] = next_address_index
self.puzzle_pk_cache[puzzle.get_tree_hash()] = sk
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
self.puzzle_pk_cache[puzzle.get_tree_hash()] = sk
self.puzzle_sk_cache[puzzle.get_tree_hash()] = sk

return puzzle

def get_new_puzzlehash(self) -> bytes32:
Expand Down