From 68263af2b55e513b102db3b4a16495b1fe7b9411 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Tue, 5 Jun 2018 17:30:16 +1200
Subject: [PATCH] Extract sapling_key_components Rust rendering

---
 sapling_key_components.py | 117 +++++++++++++-------------------------
 tv_output.py              |  15 ++++-
 2 files changed, 52 insertions(+), 80 deletions(-)

diff --git a/sapling_key_components.py b/sapling_key_components.py
index 1bae235..40ada1f 100644
--- a/sapling_key_components.py
+++ b/sapling_key_components.py
@@ -1,5 +1,4 @@
 #!/usr/bin/env python3
-from binascii import hexlify
 from pyblake2 import blake2b, blake2s
 
 from sapling_generators import PROVING_KEY_BASE, SPENDING_KEY_BASE, group_hash
@@ -7,7 +6,7 @@ from sapling_jubjub import Fr
 from sapling_merkle_tree import MERKLE_DEPTH
 from sapling_notes import note_commit, note_nullifier
 from sapling_utils import leos2bsp, leos2ip
-from tv_output import chunk
+from tv_output import tv_rust
 
 #
 # Utilities
@@ -92,26 +91,7 @@ class SpendingKey(object):
 
 
 def main():
-    print('''
-        struct TestVector {
-            sk: [u8; 32],
-            ask: [u8; 32],
-            nsk: [u8; 32],
-            ovk: [u8; 32],
-            ak: [u8; 32],
-            nk: [u8; 32],
-            ivk: [u8; 32],
-            default_d: [u8; 11],
-            default_pk_d: [u8; 32],
-            note_v: u64,
-            note_r: [u8; 32],
-            note_cm: [u8; 32],
-            note_pos: u64,
-            note_nf: [u8; 32],
-        };
-
-        // From https://github.com/zcash-hackworks/zcash-test-vectors/blob/master/sapling_key_components.py
-        let test_vectors = vec![''')
+    test_vectors = []
     for i in range(0, 10):
         sk = SpendingKey(bytes([i] * 32))
         note_v = (2548793025584392057432895043257984320*i) % 2**64
@@ -123,62 +103,43 @@ def main():
             note_v)
         note_pos = (980705743285409327583205473820957432*i) % 2**MERKLE_DEPTH
         note_nf = note_nullifier(sk.nk(), note_cm, Fr(note_pos))
-        print('''            TestVector {
-                sk: [
-                    %s
-                ],
-                ask: [
-                    %s
-                ],
-                nsk: [
-                    %s
-                ],
-                ovk: [
-                    %s
-                ],
-                ak: [
-                    %s
-                ],
-                nk: [
-                    %s
-                ],
-                ivk: [
-                    %s
-                ],
-                default_d: [
-                    %s
-                ],
-                default_pk_d: [
-                    %s
-                ],
-                note_v: %s,
-                note_r: [
-                    %s
-                ],
-                note_cm: [
-                    %s
-                ],
-                note_pos: %s,
-                note_nf: [
-                    %s
-                ],
-            },''' % (
-                chunk(hexlify(sk.data)),
-                chunk(hexlify(bytes(sk.ask()))),
-                chunk(hexlify(bytes(sk.nsk()))),
-                chunk(hexlify(sk.ovk())),
-                chunk(hexlify(bytes(sk.ak()))),
-                chunk(hexlify(bytes(sk.nk()))),
-                chunk(hexlify(bytes(sk.ivk()))),
-                chunk(hexlify(sk.default_d())),
-                chunk(hexlify(bytes(sk.default_pkd()))),
-                note_v,
-                chunk(hexlify(bytes(note_r))),
-                chunk(hexlify(bytes(note_cm.u))),
-                note_pos,
-                chunk(hexlify(note_nf)),
-            ))
-    print('        ];')
+        test_vectors.append({
+            'sk': sk.data,
+            'ask': bytes(sk.ask()),
+            'nsk': bytes(sk.nsk()),
+            'ovk': sk.ovk(),
+            'ak': bytes(sk.ak()),
+            'nk': bytes(sk.nk()),
+            'ivk': bytes(sk.ivk()),
+            'default_d': sk.default_d(),
+            'default_pk_d': bytes(sk.default_pkd()),
+            'note_v': note_v,
+            'note_r': bytes(note_r),
+            'note_cm': bytes(note_cm.u),
+            'note_pos': note_pos,
+            'note_nf': note_nf,
+        })
+
+    tv_rust(
+        'sapling_key_components',
+        (
+            ('sk', '[u8; 32]'),
+            ('ask', '[u8; 32]'),
+            ('nsk', '[u8; 32]'),
+            ('ovk', '[u8; 32]'),
+            ('ak', '[u8; 32]'),
+            ('nk', '[u8; 32]'),
+            ('ivk', '[u8; 32]'),
+            ('default_d', '[u8; 11]'),
+            ('default_pk_d', '[u8; 32]'),
+            ('note_v', 'u64'),
+            ('note_r', '[u8; 32]'),
+            ('note_cm', '[u8; 32]'),
+            ('note_pos', 'u64'),
+            ('note_nf', '[u8; 32]'),
+        ),
+        test_vectors,
+    )
 
 
 if __name__ == '__main__':
diff --git a/tv_output.py b/tv_output.py
index fae030b..53c09b7 100644
--- a/tv_output.py
+++ b/tv_output.py
@@ -5,8 +5,7 @@ def chunk(h):
     h = str(h, 'utf-8')
     return '0x' + ', 0x'.join([h[i:i+2] for i in range(0, len(h), 2)])
 
-def tv_part_rust(name, value, indent=3):
-    pad = '    ' * indent
+def tv_bytes_rust(name, value, pad):
     print('''%s%s: [
     %s%s
 %s],''' % (
@@ -17,6 +16,18 @@ def tv_part_rust(name, value, indent=3):
         pad,
     ))
 
+def tv_int_rust(name, value, pad):
+    print('%s%s: %d,' % (pad, name, value))
+
+def tv_part_rust(name, value, indent=3):
+    pad = '    ' * indent
+    if type(value) == bytes:
+        tv_bytes_rust(name, value, pad)
+    elif type(value) == int:
+        tv_int_rust(name, value, pad)
+    else:
+        raise ValueError('Invalid type(%s): %s' % (name, type(value)))
+
 def tv_rust(filename, parts, vectors):
     print('        struct TestVector {')
     [print('            %s: %s,' % p) for p in parts]