From c1b9e9ca1407ed90f0fa55d0843f66017bcad964 Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Wed, 6 Aug 2025 17:19:38 +0000
Subject: [PATCH] chore: make `TokenFlags` fields booleans

---
 src/_table_generation/make_tables.nr | 138 +++++++++++++--------------
 src/json.nr                          |   9 ++
 src/token_flags.nr                   |  69 ++++----------
 3 files changed, 99 insertions(+), 117 deletions(-)

diff --git a/src/_table_generation/make_tables.nr b/src/_table_generation/make_tables.nr
index f3a7e6f..65c04fb 100644
--- a/src/_table_generation/make_tables.nr
+++ b/src/_table_generation/make_tables.nr
@@ -359,94 +359,94 @@ unconstrained fn make_process_raw_transcript_table() -> [Field; 1024] {
 }
 
 unconstrained fn generate_token_flags_table() -> [Field; NUM_TOKENS * 2] {
-    let mut flags: [TokenFlags; NUM_TOKENS * 2] = [TokenFlags::default(); NUM_TOKENS * 2];
+    let mut flags: [TokenFlags; NUM_TOKENS * 2] = std::mem::zeroed();
 
     let mut no_token_flags: TokenFlags = TokenFlags {
-        create_json_entry: 0,
-        is_end_of_object_or_array: 0,
-        is_start_of_object_or_array: 0,
-        new_context: OBJECT_LAYER as Field,
-        is_key_token: 0,
-        is_value_token: 0,
-        preserve_num_entries: 1,
+        create_json_entry: false,
+        is_end_of_object_or_array: false,
+        is_start_of_object_or_array: false,
+        new_context: OBJECT_LAYER != 0,
+        is_key_token: false,
+        is_value_token: false,
+        preserve_num_entries: true,
     };
     let mut key_token_flags: TokenFlags = TokenFlags {
-        create_json_entry: 0,
-        is_end_of_object_or_array: 0,
-        is_start_of_object_or_array: 0,
-        new_context: OBJECT_LAYER as Field,
-        is_key_token: 1,
-        is_value_token: 0,
-        preserve_num_entries: 1,
+        create_json_entry: false,
+        is_end_of_object_or_array: false,
+        is_start_of_object_or_array: false,
+        new_context: OBJECT_LAYER != 0,
+        is_key_token: true,
+        is_value_token: false,
+        preserve_num_entries: true,
     };
     let begin_object_flags = TokenFlags {
-        create_json_entry: 0,
-        is_end_of_object_or_array: 0,
-        is_start_of_object_or_array: 1,
-        new_context: OBJECT_LAYER as Field,
-        is_key_token: 0,
-        is_value_token: 0,
-        preserve_num_entries: 0,
+        create_json_entry: false,
+        is_end_of_object_or_array: false,
+        is_start_of_object_or_array: true,
+        new_context: OBJECT_LAYER != 0,
+        is_key_token: false,
+        is_value_token: false,
+        preserve_num_entries: false,
     };
 
     let begin_array_flags = TokenFlags {
-        create_json_entry: 0,
-        is_end_of_object_or_array: 0,
-        is_start_of_object_or_array: 1,
-        new_context: ARRAY_LAYER as Field,
-        is_key_token: 0,
-        is_value_token: 0,
-        preserve_num_entries: 0,
+        create_json_entry: false,
+        is_end_of_object_or_array: false,
+        is_start_of_object_or_array: true,
+        new_context: ARRAY_LAYER != 0,
+        is_key_token: false,
+        is_value_token: false,
+        preserve_num_entries: false,
     };
 
     let mut end_object_flags = TokenFlags {
-        create_json_entry: 1,
-        is_end_of_object_or_array: 1,
-        is_start_of_object_or_array: 0,
-        new_context: 0,
-        is_key_token: 0,
-        is_value_token: 0,
-        preserve_num_entries: 0,
+        create_json_entry: true,
+        is_end_of_object_or_array: true,
+        is_start_of_object_or_array: false,
+        new_context: false,
+        is_key_token: false,
+        is_value_token: false,
+        preserve_num_entries: false,
     };
 
     let mut end_array_flags = TokenFlags {
-        create_json_entry: 1,
-        is_end_of_object_or_array: 1,
-        is_start_of_object_or_array: 0,
-        new_context: 0,
-        is_key_token: 0,
-        is_value_token: 0,
-        preserve_num_entries: 0,
+        create_json_entry: true,
+        is_end_of_object_or_array: true,
+        is_start_of_object_or_array: false,
+        new_context: false,
+        is_key_token: false,
+        is_value_token: false,
+        preserve_num_entries: false,
     };
 
     let mut string_flags = TokenFlags {
-        create_json_entry: 1,
-        is_end_of_object_or_array: 0,
-        is_start_of_object_or_array: 0,
-        new_context: OBJECT_LAYER as Field,
-        is_key_token: 0,
-        is_value_token: 1,
-        preserve_num_entries: 1,
+        create_json_entry: true,
+        is_end_of_object_or_array: false,
+        is_start_of_object_or_array: false,
+        new_context: OBJECT_LAYER != 0,
+        is_key_token: false,
+        is_value_token: true,
+        preserve_num_entries: true,
     };
 
     let mut numeric_flags = TokenFlags {
-        create_json_entry: 1,
-        is_end_of_object_or_array: 0,
-        is_start_of_object_or_array: 0,
-        new_context: OBJECT_LAYER as Field,
-        is_key_token: 0,
-        is_value_token: 1,
-        preserve_num_entries: 1,
+        create_json_entry: true,
+        is_end_of_object_or_array: false,
+        is_start_of_object_or_array: false,
+        new_context: OBJECT_LAYER != 0,
+        is_key_token: false,
+        is_value_token: true,
+        preserve_num_entries: true,
     };
 
     let mut literal_flags = TokenFlags {
-        create_json_entry: 1,
-        is_end_of_object_or_array: 0,
-        is_start_of_object_or_array: 0,
-        new_context: OBJECT_LAYER as Field,
-        is_key_token: 0,
-        is_value_token: 1,
-        preserve_num_entries: 1,
+        create_json_entry: true,
+        is_end_of_object_or_array: false,
+        is_start_of_object_or_array: false,
+        new_context: OBJECT_LAYER != 0,
+        is_key_token: false,
+        is_value_token: true,
+        preserve_num_entries: true,
     };
 
     flags[NO_TOKEN] = no_token_flags;
@@ -461,11 +461,11 @@ unconstrained fn generate_token_flags_table() -> [Field; NUM_TOKENS * 2] {
     flags[LITERAL_TOKEN] = literal_flags;
     flags[KEY_TOKEN] = key_token_flags;
 
-    no_token_flags.new_context = ARRAY_LAYER as Field;
-    key_token_flags.new_context = ARRAY_LAYER as Field;
-    string_flags.new_context = ARRAY_LAYER as Field;
-    numeric_flags.new_context = ARRAY_LAYER as Field;
-    literal_flags.new_context = ARRAY_LAYER as Field;
+    no_token_flags.new_context = ARRAY_LAYER != 0;
+    key_token_flags.new_context = ARRAY_LAYER != 0;
+    string_flags.new_context = ARRAY_LAYER != 0;
+    numeric_flags.new_context = ARRAY_LAYER != 0;
+    literal_flags.new_context = ARRAY_LAYER != 0;
 
     flags[NUM_TOKENS + (NO_TOKEN as u32)] = no_token_flags;
     flags[NUM_TOKENS + (BEGIN_OBJECT_TOKEN as u32)] = begin_object_flags;
diff --git a/src/json.nr b/src/json.nr
index e362525..cde56a8 100644
--- a/src/json.nr
+++ b/src/json.nr
@@ -276,6 +276,15 @@ impl<let NumBytes: u32, let NumPackedFields: u32, let MaxNumTokens: u32, let Max
                 TOKEN_FLAGS_TABLE[cast_num_to_u32(token) + context * NUM_TOKENS],
             );
 
+            // We convert these booleans into Fields so that we can use them in arithmetic operations.
+            let create_json_entry = create_json_entry as Field;
+            let is_end_of_object_or_array = is_end_of_object_or_array as Field;
+            let is_start_of_object_or_array = is_start_of_object_or_array as Field;
+            let new_context = new_context as Field;
+            let update_key = update_key as Field;
+            let is_value_token = is_value_token as Field;
+            let preserve_num_entries = preserve_num_entries as Field;
+
             // 2 gates
             let diff = (index + length * 0x10000) - current_key_index_and_length;
             std::as_witness(diff);
diff --git a/src/token_flags.nr b/src/token_flags.nr
index abbf2cb..da9aa79 100644
--- a/src/token_flags.nr
+++ b/src/token_flags.nr
@@ -1,24 +1,24 @@
 pub(crate) struct TokenFlags {
-    pub(crate) create_json_entry: Field,
-    pub(crate) is_end_of_object_or_array: Field,
-    pub(crate) is_start_of_object_or_array: Field,
-    pub(crate) new_context: Field,
-    pub(crate) is_key_token: Field,
-    pub(crate) is_value_token: Field,
-    pub(crate) preserve_num_entries: Field,
+    pub(crate) create_json_entry: bool,
+    pub(crate) is_end_of_object_or_array: bool,
+    pub(crate) is_start_of_object_or_array: bool,
+    pub(crate) new_context: bool,
+    pub(crate) is_key_token: bool,
+    pub(crate) is_value_token: bool,
+    pub(crate) preserve_num_entries: bool,
 }
 
 impl TokenFlags {
 
     unconstrained fn __from_field(f: Field) -> Self {
         let bytes: [u8; 7] = f.to_be_bytes();
-        let create_json_entry = bytes[0] as Field;
-        let is_end_of_object_or_array = bytes[1] as Field;
-        let is_start_of_object_or_array = bytes[2] as Field;
-        let new_context = bytes[3] as Field;
-        let is_key_token = bytes[4] as Field;
-        let is_value_token = bytes[5] as Field;
-        let preserve_num_entries = bytes[6] as Field;
+        let create_json_entry = bytes[0] != 0;
+        let is_end_of_object_or_array = bytes[1] != 0;
+        let is_start_of_object_or_array = bytes[2] != 0;
+        let new_context = bytes[3] != 0;
+        let is_key_token = bytes[4] != 0;
+        let is_value_token = bytes[5] != 0;
+        let preserve_num_entries = bytes[6] != 0;
 
         TokenFlags {
             create_json_entry,
@@ -36,21 +36,6 @@ impl TokenFlags {
         // Safety: check the comments below
         let r = unsafe { TokenFlags::__from_field(f) };
 
-        // checks that the flags are binary
-        assert(r.create_json_entry * r.create_json_entry == r.create_json_entry);
-        assert(
-            r.is_end_of_object_or_array * r.is_end_of_object_or_array
-                == r.is_end_of_object_or_array,
-        );
-        assert(
-            r.is_start_of_object_or_array * r.is_start_of_object_or_array
-                == r.is_start_of_object_or_array,
-        );
-        assert(r.new_context * r.new_context == r.new_context);
-        assert(r.is_key_token * r.is_key_token == r.is_key_token);
-        assert(r.is_value_token * r.is_value_token == r.is_value_token);
-        assert(r.preserve_num_entries * r.preserve_num_entries == r.preserve_num_entries);
-
         // asserts the relation of r and f
         assert(r.to_field() == f);
         r
@@ -58,24 +43,12 @@ impl TokenFlags {
 
     // 4 gates
     pub(crate) fn to_field(self) -> Field {
-        self.preserve_num_entries
-            + self.is_value_token * 0x100
-            + self.is_key_token * 0x10000
-            + self.new_context * 0x1000000
-            + self.is_start_of_object_or_array * 0x100000000
-            + self.is_end_of_object_or_array * 0x10000000000
-            + self.create_json_entry * 0x1000000000000
-    }
-
-    pub(crate) fn default() -> Self {
-        TokenFlags {
-            create_json_entry: 0,
-            is_end_of_object_or_array: 0,
-            is_start_of_object_or_array: 0,
-            new_context: 0,
-            is_key_token: 0,
-            preserve_num_entries: 0,
-            is_value_token: 0,
-        }
+        (self.preserve_num_entries as Field)
+            + (self.is_value_token as Field) * 0x100
+            + (self.is_key_token as Field) * 0x10000
+            + (self.new_context as Field) * 0x1000000
+            + (self.is_start_of_object_or_array as Field) * 0x100000000
+            + (self.is_end_of_object_or_array as Field) * 0x10000000000
+            + (self.create_json_entry as Field) * 0x1000000000000
     }
 }