Better conversion functions

argumentcomputer · Jan 5, 2024 · def5463 · def5463
1 parent 337ef0a
commit def5463
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 119 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -53,6 +53,7 @@ serde = { workspace = true, features = ["derive"] }
 serde_bytes = "0.11.12"
 serde_json = { workspace = true }
 serde_repr = "0.1.14"
+strum = { version = "0.25", features = ["derive"] }
 tap = "1.0.1"
 stable_deref_trait = "1.2.0"
 thiserror = { workspace = true }

diff --git a/src/lem/tag.rs b/src/lem/tag.rs
@@ -1,9 +1,13 @@
 use anyhow::{bail, Result};
 use serde::{Deserialize, Serialize};
+use strum::EnumCount;
 
 use crate::{
     field::LurkField,
-    tag::{ContTag, ExprTag, Op1, Op2, Tag as TagTrait},
+    tag::{
+        ContTag, ExprTag, Op1, Op2, Tag as TagTrait, CONT_TAG_INIT, EXPR_TAG_INIT, OP1_TAG_INIT,
+        OP2_TAG_INIT,
+    },
 };
 
 /// The LEM `Tag` is a wrapper around other types that are used as tags
@@ -66,127 +70,43 @@ impl Tag {
     }
 
     pub fn pos(i: usize) -> Option<Self> {
-        match i {
-            0 => Some(Self::Expr(ExprTag::Nil)),
-            1 => Some(Self::Expr(ExprTag::Cons)),
-            2 => Some(Self::Expr(ExprTag::Sym)),
-            3 => Some(Self::Expr(ExprTag::Fun)),
-            4 => Some(Self::Expr(ExprTag::Num)),
-            5 => Some(Self::Expr(ExprTag::Thunk)),
-            6 => Some(Self::Expr(ExprTag::Str)),
-            7 => Some(Self::Expr(ExprTag::Char)),
-            8 => Some(Self::Expr(ExprTag::Comm)),
-            9 => Some(Self::Expr(ExprTag::U64)),
-            10 => Some(Self::Expr(ExprTag::Key)),
-            11 => Some(Self::Expr(ExprTag::Cproc)),
-            12 => Some(Self::Cont(ContTag::Outermost)),
-            13 => Some(Self::Cont(ContTag::Call0)),
-            14 => Some(Self::Cont(ContTag::Call)),
-            15 => Some(Self::Cont(ContTag::Call2)),
-            16 => Some(Self::Cont(ContTag::Tail)),
-            17 => Some(Self::Cont(ContTag::Error)),
-            18 => Some(Self::Cont(ContTag::Lookup)),
-            19 => Some(Self::Cont(ContTag::Unop)),
-            20 => Some(Self::Cont(ContTag::Binop)),
-            21 => Some(Self::Cont(ContTag::Binop2)),
-            22 => Some(Self::Cont(ContTag::If)),
-            23 => Some(Self::Cont(ContTag::Let)),
-            24 => Some(Self::Cont(ContTag::LetRec)),
-            25 => Some(Self::Cont(ContTag::Dummy)),
-            26 => Some(Self::Cont(ContTag::Terminal)),
-            27 => Some(Self::Cont(ContTag::Emit)),
-            28 => Some(Self::Cont(ContTag::Cproc)),
-            29 => Some(Self::Op1(Op1::Car)),
-            30 => Some(Self::Op1(Op1::Cdr)),
-            31 => Some(Self::Op1(Op1::Atom)),
-            32 => Some(Self::Op1(Op1::Emit)),
-            33 => Some(Self::Op1(Op1::Open)),
-            34 => Some(Self::Op1(Op1::Secret)),
-            35 => Some(Self::Op1(Op1::Commit)),
-            36 => Some(Self::Op1(Op1::Num)),
-            37 => Some(Self::Op1(Op1::Comm)),
-            38 => Some(Self::Op1(Op1::Char)),
-            39 => Some(Self::Op1(Op1::Eval)),
-            40 => Some(Self::Op1(Op1::U64)),
-            41 => Some(Self::Op2(Op2::Sum)),
-            42 => Some(Self::Op2(Op2::Diff)),
-            43 => Some(Self::Op2(Op2::Product)),
-            44 => Some(Self::Op2(Op2::Quotient)),
-            45 => Some(Self::Op2(Op2::Equal)),
-            46 => Some(Self::Op2(Op2::NumEqual)),
-            47 => Some(Self::Op2(Op2::Less)),
-            48 => Some(Self::Op2(Op2::Greater)),
-            49 => Some(Self::Op2(Op2::LessEqual)),
-            50 => Some(Self::Op2(Op2::GreaterEqual)),
-            51 => Some(Self::Op2(Op2::Cons)),
-            52 => Some(Self::Op2(Op2::StrCons)),
-            53 => Some(Self::Op2(Op2::Begin)),
-            54 => Some(Self::Op2(Op2::Hide)),
-            55 => Some(Self::Op2(Op2::Modulo)),
-            56 => Some(Self::Op2(Op2::Eval)),
-            _ => None,
+        let mut last = 0;
+        if (last..last + ExprTag::COUNT).contains(&i) {
+            let j = i + EXPR_TAG_INIT as usize - last;
+            let expr_tag = (j as u16).try_into().expect("unreachable");
+            return Some(Tag::Expr(expr_tag));
         }
+        last += ExprTag::COUNT;
+        if (last..last + ContTag::COUNT).contains(&i) {
+            let j = i + CONT_TAG_INIT as usize - last;
+            let cont_tag = (j as u16).try_into().expect("unreachable");
+            return Some(Tag::Cont(cont_tag));
+        }
+        last += ContTag::COUNT;
+        if (last..last + Op1::COUNT).contains(&i) {
+            let j = i + OP1_TAG_INIT as usize - last;
+            let op1_tag = (j as u16).try_into().expect("unreachable");
+            return Some(Tag::Op1(op1_tag));
+        }
+        last += Op1::COUNT;
+        if (last..last + Op2::COUNT).contains(&i) {
+            let j = i + OP2_TAG_INIT as usize - last;
+            let op2_tag = (j as u16).try_into().expect("unreachable");
+            return Some(Tag::Op2(op2_tag));
+        }
+        None
     }
 
     pub fn index(&self) -> usize {
         match self {
-            Self::Expr(ExprTag::Nil) => 0,
-            Self::Expr(ExprTag::Cons) => 1,
-            Self::Expr(ExprTag::Sym) => 2,
-            Self::Expr(ExprTag::Fun) => 3,
-            Self::Expr(ExprTag::Num) => 4,
-            Self::Expr(ExprTag::Thunk) => 5,
-            Self::Expr(ExprTag::Str) => 6,
-            Self::Expr(ExprTag::Char) => 7,
-            Self::Expr(ExprTag::Comm) => 8,
-            Self::Expr(ExprTag::U64) => 9,
-            Self::Expr(ExprTag::Key) => 10,
-            Self::Expr(ExprTag::Cproc) => 11,
-            Self::Cont(ContTag::Outermost) => 12,
-            Self::Cont(ContTag::Call0) => 13,
-            Self::Cont(ContTag::Call) => 14,
-            Self::Cont(ContTag::Call2) => 15,
-            Self::Cont(ContTag::Tail) => 16,
-            Self::Cont(ContTag::Error) => 17,
-            Self::Cont(ContTag::Lookup) => 18,
-            Self::Cont(ContTag::Unop) => 19,
-            Self::Cont(ContTag::Binop) => 20,
-            Self::Cont(ContTag::Binop2) => 21,
-            Self::Cont(ContTag::If) => 22,
-            Self::Cont(ContTag::Let) => 23,
-            Self::Cont(ContTag::LetRec) => 24,
-            Self::Cont(ContTag::Dummy) => 25,
-            Self::Cont(ContTag::Terminal) => 26,
-            Self::Cont(ContTag::Emit) => 27,
-            Self::Cont(ContTag::Cproc) => 28,
-            Self::Op1(Op1::Car) => 29,
-            Self::Op1(Op1::Cdr) => 30,
-            Self::Op1(Op1::Atom) => 31,
-            Self::Op1(Op1::Emit) => 32,
-            Self::Op1(Op1::Open) => 33,
-            Self::Op1(Op1::Secret) => 34,
-            Self::Op1(Op1::Commit) => 35,
-            Self::Op1(Op1::Num) => 36,
-            Self::Op1(Op1::Comm) => 37,
-            Self::Op1(Op1::Char) => 38,
-            Self::Op1(Op1::Eval) => 39,
-            Self::Op1(Op1::U64) => 40,
-            Self::Op2(Op2::Sum) => 41,
-            Self::Op2(Op2::Diff) => 42,
-            Self::Op2(Op2::Product) => 43,
-            Self::Op2(Op2::Quotient) => 44,
-            Self::Op2(Op2::Equal) => 45,
-            Self::Op2(Op2::NumEqual) => 46,
-            Self::Op2(Op2::Less) => 47,
-            Self::Op2(Op2::Greater) => 48,
-            Self::Op2(Op2::LessEqual) => 49,
-            Self::Op2(Op2::GreaterEqual) => 50,
-            Self::Op2(Op2::Cons) => 51,
-            Self::Op2(Op2::StrCons) => 52,
-            Self::Op2(Op2::Begin) => 53,
-            Self::Op2(Op2::Hide) => 54,
-            Self::Op2(Op2::Modulo) => 55,
-            Self::Op2(Op2::Eval) => 56,
+            Self::Expr(tag) => *tag as usize - EXPR_TAG_INIT as usize,
+            Self::Cont(tag) => *tag as usize - CONT_TAG_INIT as usize + ExprTag::COUNT,
+            Self::Op1(tag) => {
+                *tag as usize - OP1_TAG_INIT as usize + ExprTag::COUNT + ContTag::COUNT
+            }
+            Self::Op2(tag) => {
+                *tag as usize - OP2_TAG_INIT as usize + ExprTag::COUNT + ContTag::COUNT + Op1::COUNT
+            }
         }
     }
 }

diff --git a/src/tag.rs b/src/tag.rs
@@ -3,6 +3,7 @@ use lurk_macros::TryFromRepr;
 use proptest_derive::Arbitrary;
 use serde_repr::{Deserialize_repr, Serialize_repr};
 use std::{convert::TryFrom, fmt};
+use strum::EnumCount;
 
 use crate::field::LurkField;
 
@@ -23,7 +24,16 @@ pub trait Tag:
 pub(crate) const EXPR_TAG_INIT: u16 = 0b0000_0000_0000_0000;
 /// A tag for expressions. Note that ExprTag, ContTag, Op1, Op2 all live in the same u16 namespace
 #[derive(
-    Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize_repr, Deserialize_repr, TryFromRepr,
+    Debug,
+    Copy,
+    Clone,
+    PartialEq,
+    Eq,
+    Hash,
+    Serialize_repr,
+    Deserialize_repr,
+    TryFromRepr,
+    EnumCount,
 )]
 #[cfg_attr(not(target_arch = "wasm32"), derive(Arbitrary))]
 #[repr(u16)]
@@ -93,7 +103,16 @@ impl Tag for ExprTag {
 pub(crate) const CONT_TAG_INIT: u16 = 0b0001_0000_0000_0000;
 /// A tag for continuations. Note that ExprTag, ContTag, Op1, Op2 all live in the same u16 namespace
 #[derive(
-    Serialize_repr, Deserialize_repr, Debug, Copy, Clone, PartialEq, Eq, Hash, TryFromRepr,
+    Serialize_repr,
+    Deserialize_repr,
+    Debug,
+    Copy,
+    Clone,
+    PartialEq,
+    Eq,
+    Hash,
+    TryFromRepr,
+    EnumCount,
 )]
 #[cfg_attr(not(target_arch = "wasm32"), derive(Arbitrary))]
 #[repr(u16)]
@@ -182,6 +201,7 @@ pub(crate) const OP1_TAG_INIT: u16 = 0b0010_0000_0000_0000;
     Serialize_repr,
     Deserialize_repr,
     TryFromRepr,
+    EnumCount,
 )]
 #[cfg_attr(not(target_arch = "wasm32"), derive(Arbitrary))]
 #[repr(u16)]
@@ -312,6 +332,7 @@ pub(crate) const OP2_TAG_INIT: u16 = 0b0011_0000_0000_0000;
     Serialize_repr,
     Deserialize_repr,
     TryFromRepr,
+    EnumCount,
 )]
 #[cfg_attr(not(target_arch = "wasm32"), derive(Arbitrary))]
 #[repr(u16)]