implement enum (#15)

Author: jackdotink

src/api.rs

@@ -125,6 +125,7 @@ fn library() -> lu::Library<Config> {
         .with_function_norm("array", array)
         .with_function_norm("set", set)
         .with_function_norm("map", map)
+        .with_function_norm("enum", enumn)
         .with_function_norm("struct", strukt)
 }
 
@@ -175,6 +176,7 @@ pub enum Type {
     Array(ArrayType),
     Set(SetType),
     Map(MapType),
+    Enum(EnumType),
     Struct(StructType),
 }
 
@@ -673,6 +675,28 @@ extern "C-unwind" fn map(ctx: Context) -> lu::FnReturn {
     ctx.ret_with(1)
 }
 
+#[derive(Clone)]
+pub struct EnumType {
+    pub variants: Vec<String>,
+}
+
+extern "C-unwind" fn enumn(ctx: Context) -> lu::FnReturn {
+    ctx.arg_table(1);
+    let mut variants = Vec::new();
+
+    ctx.iter(1, || {
+        let Some(str) = ctx.to_string_str(-1) else {
+            ctx.error_msg("enum variant must be a string");
+        };
+
+        variants.push(str.to_string());
+        ControlFlow::<()>::Continue(())
+    });
+
+    ctx.push_userdata(Type::Enum(EnumType { variants }));
+    ctx.ret_with(1)
+}
+
 #[derive(Clone)]
 pub struct StructType {
     pub fields: Vec<(String, Type)>,

src/hir/build.rs

@@ -1,8 +1,8 @@
 use crate::{
     api,
     hir::{
-        ArrayType, BinaryStringType, Event, Item, Length, MapType, NumberType, SetType, StructType,
-        Type, Utf8StringType, VectorType,
+        ArrayType, BinaryStringType, EnumType, Event, Item, Length, MapType, NumberType, SetType,
+        StructType, Type, Utf8StringType, VectorType,
     },
     shared::Range,
 };
@@ -42,6 +42,7 @@ impl From<api::Type> for Type {
             api::Type::Array(ty) => Type::Array(ty.into()),
             api::Type::Set(ty) => Type::Set(ty.into()),
             api::Type::Map(ty) => Type::Map(ty.into()),
+            api::Type::Enum(ty) => Type::Enum(ty.into()),
             api::Type::Struct(ty) => Type::Struct(ty.into()),
         }
     }
@@ -110,6 +111,22 @@ impl From<api::MapType> for MapType {
     }
 }
 
+impl From<api::EnumType> for EnumType {
+    fn from(value: api::EnumType) -> Self {
+        let variants = value.variants;
+        let range = Range {
+            min: Some(0f64),
+            max: Some(variants.len() as f64 - 1f64),
+        };
+        let number = NumberType {
+            kind: range.kind(),
+            range,
+        };
+
+        EnumType { variants, number }
+    }
+}
+
 impl From<api::StructType> for StructType {
     fn from(value: api::StructType) -> Self {
         let fields = value

src/hir/mod.rs

@@ -27,6 +27,7 @@ pub enum Type {
     Array(ArrayType),
     Set(SetType),
     Map(MapType),
+    Enum(EnumType),
     Struct(StructType),
 }
 
@@ -72,6 +73,12 @@ pub struct MapType {
     pub value: Box<Type>,
 }
 
+#[derive(Clone)]
+pub struct EnumType {
+    pub variants: Vec<String>,
+    pub number: NumberType,
+}
+
 #[derive(Clone)]
 pub struct StructType {
     pub fields: Vec<(String, Type)>,

src/hir/size.rs

@@ -90,6 +90,7 @@ impl hir::Type {
             hir::Type::Array(ty) => ty.size(),
             hir::Type::Set(ty) => ty.size(),
             hir::Type::Map(ty) => ty.size(),
+            hir::Type::Enum(ty) => ty.size(),
             hir::Type::Struct(ty) => ty.size(),
         }
     }
@@ -140,6 +141,12 @@ impl hir::MapType {
     }
 }
 
+impl hir::EnumType {
+    pub fn size(&self) -> Size {
+        self.number.size()
+    }
+}
+
 impl hir::StructType {
     pub fn size(&self) -> Size {
         self.fields

src/mir/serdes.rs

@@ -125,6 +125,7 @@ impl Serdes for hir::Type {
             hir::Type::Array(ty) => Box::new(ty.ser(b, ser)),
             hir::Type::Set(ty) => Box::new(ty.ser(b, ser)),
             hir::Type::Map(ty) => Box::new(ty.ser(b, ser)),
+            hir::Type::Enum(ty) => Box::new(ty.ser(b, ser)),
             hir::Type::Struct(ty) => Box::new(ty.ser(b, ser)),
         };
 
@@ -146,6 +147,7 @@ impl Serdes for hir::Type {
             hir::Type::Array(ty) => Box::new(ty.des(b, des)),
             hir::Type::Set(ty) => Box::new(ty.des(b, des)),
             hir::Type::Map(ty) => Box::new(ty.des(b, des)),
+            hir::Type::Enum(ty) => Box::new(ty.des(b, des)),
             hir::Type::Struct(ty) => Box::new(ty.des(b, des)),
         };
 
@@ -439,6 +441,54 @@ impl Serdes for hir::MapType {
     }
 }
 
+impl Serdes for hir::EnumType {
+    fn ser<'ty, 'b, 'ser: 'ty>(
+        &'ty self,
+        b: &'b mut Builder,
+        ser: &'ser Ser,
+    ) -> impl Fn(&mut Builder, Expr) + use<'ty, 'ser> + 'ty {
+        let variants = b.expr(Expr::Table(
+            self.variants
+                .iter()
+                .enumerate()
+                .map(|(i, v)| (Expr::from(v.as_str()), Expr::from(i as f64)))
+                .collect(),
+        ));
+
+        let number = self.number.ser(b, ser);
+
+        move |b: &mut Builder, from: Expr| {
+            apicheck_full!(ser, check_type(b, from.clone(), "string"));
+
+            let value = b.expr(variants.expr().index(from));
+            apicheck_some!(ser, b.assert(value.expr(), "not a valid enum variant"));
+
+            number(b, value.expr());
+        }
+    }
+
+    fn des<'ty, 'b, 'des: 'ty>(
+        &'ty self,
+        b: &'b mut Builder,
+        des: &'des Des,
+    ) -> impl Fn(&mut Builder) -> InitVar + use<'ty, 'des> + 'ty {
+        let variants = b.expr(Expr::Table(
+            self.variants
+                .iter()
+                .enumerate()
+                .map(|(i, v)| (Expr::from(i as f64), Expr::from(v.as_str())))
+                .collect(),
+        ));
+
+        let number = self.number.des(b, des);
+
+        move |b: &mut Builder| {
+            let value = number(b);
+            b.expr(variants.expr().index(&value))
+        }
+    }
+}
+
 impl Serdes for hir::StructType {
     fn ser<'ty, 'b, 'ser: 'ty>(
         &'ty self,

src/shared.rs

@@ -71,7 +71,7 @@ impl Range {
         if self.min == self.max { self.min } else { None }
     }
 
-    pub fn len_kind(&self) -> NumberKind {
+    pub fn kind(&self) -> NumberKind {
         let max = self.max.unwrap_or(f64::MAX);
 
         if max <= u8::MAX as f64 {