118 lines
3.3 KiB
Nim
118 lines
3.3 KiB
Nim
import
|
|
randutils, random, parseopt, strutils, os,
|
|
eth/[common, rlp], eth/trie/[hexary, db, trie_defs],
|
|
nimcrypto/sysrand, ../stateless/[json_from_tree],
|
|
../nimbus/db/storage_types, ./witness_types, ./multi_keys
|
|
|
|
type
|
|
DB = TrieDatabaseRef
|
|
|
|
StorageKeys = tuple[storageRoot: Hash256, keys: MultikeysRef]
|
|
|
|
AccountDef = object
|
|
storageKeys: MultiKeysRef
|
|
account: Account
|
|
codeTouched: bool
|
|
|
|
proc randU256(): UInt256 =
|
|
var bytes: array[32, byte]
|
|
discard randomBytes(bytes[0].addr, sizeof(result))
|
|
result = UInt256.fromBytesBE(bytes)
|
|
|
|
proc randStorageSlot(): StorageSlot =
|
|
discard randomBytes(result[0].addr, sizeof(result))
|
|
|
|
proc randNonce(): AccountNonce =
|
|
discard randomBytes(result.addr, sizeof(result))
|
|
|
|
proc randCode(db: DB): Hash256 =
|
|
if rand(0..1) == 0:
|
|
result = blankStringHash
|
|
else:
|
|
let codeLen = rand(1..150)
|
|
let code = randList(byte, rng(0, 255), codeLen, unique = false)
|
|
result = keccakHash(code)
|
|
db.put(contractHashKey(result).toOpenArray, code)
|
|
|
|
proc randStorage(db: DB, numSlots: int): StorageKeys =
|
|
if rand(0..1) == 0 or numSlots == 0:
|
|
result = (emptyRlpHash, MultikeysRef(nil))
|
|
else:
|
|
var trie = initSecureHexaryTrie(db)
|
|
var keys = newSeq[StorageSlot](numSlots)
|
|
|
|
for i in 0..<numSlots:
|
|
keys[i] = randStorageSlot()
|
|
trie.put(keys[i], rlp.encode(randU256()))
|
|
|
|
if rand(0..1) == 0:
|
|
result = (trie.rootHash, MultikeysRef(nil))
|
|
else:
|
|
var m = newMultikeys(keys)
|
|
result = (trie.rootHash, m)
|
|
|
|
proc randAccount(db: DB, numSlots: int): AccountDef =
|
|
result.account.nonce = randNonce()
|
|
result.account.balance = randU256()
|
|
let z = randStorage(db, numSlots)
|
|
result.account.codeHash = randCode(db)
|
|
result.account.storageRoot = z.storageRoot
|
|
result.storageKeys = z.keys
|
|
result.codeTouched = rand(0..1) == 0
|
|
|
|
proc randAddress(): EthAddress =
|
|
discard randomBytes(result.addr, sizeof(result))
|
|
|
|
proc runGenerator(numPairs, numSlots: int): string =
|
|
var memDB = newMemoryDB()
|
|
var trie = initSecureHexaryTrie(memDB)
|
|
var addrs = newSeq[AccountKey](numPairs)
|
|
var accs = newSeq[Account](numPairs)
|
|
|
|
for i in 0..<numPairs:
|
|
let acc = randAccount(memDB, numSlots)
|
|
addrs[i] = (randAddress(), acc.codeTouched, acc.storageKeys)
|
|
accs[i] = acc.account
|
|
trie.put(addrs[i].address, rlp.encode(accs[i]))
|
|
|
|
var mkeys = newMultiKeys(addrs)
|
|
let rootHash = trie.rootHash
|
|
|
|
var wb = initWitnessBuilder(memDB, rootHash, {wfEIP170})
|
|
result = wb.buildWitness(mkeys)
|
|
|
|
proc writeHelp() =
|
|
echo "json_witness_gen output --pairs:val --slots:val -s:val -p:val"
|
|
|
|
proc main() =
|
|
var filename: string
|
|
var outputDir: string
|
|
var numPairs = 1
|
|
var numSlots = 1
|
|
for kind, key, val in getopt():
|
|
case kind
|
|
of cmdArgument:
|
|
filename = key
|
|
of cmdLongOption, cmdShortOption:
|
|
case key
|
|
of "pairs", "p":
|
|
numPairs = parseInt(val)
|
|
if numPairs <= 0: numPairs = 1
|
|
of "slots", "s":
|
|
numSlots = parseInt(val)
|
|
if numSlots < 0: numSlots = 0
|
|
of "output", "o":
|
|
outputDir = val
|
|
of cmdEnd: assert(false) # cannot happen
|
|
|
|
if filename == "":
|
|
writeHelp()
|
|
quit(0)
|
|
|
|
randomize()
|
|
let witness = runGenerator(numPairs, numSlots)
|
|
let filePath = if outputDir.len > 0: outputDir / filename: else: filename
|
|
writeFile(filePath, witness)
|
|
|
|
main()
|