From 087d926c3d00cf0a53e43688f867808fba6224e5 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Mon, 18 Mar 2024 16:52:09 -0400
Subject: [PATCH] Implement enums

---
 src/abstract_tree/enum_definition.rs |  8 +++----
 src/abstract_tree/type.rs            |  4 ++--
 src/abstract_tree/value_node.rs      | 34 ++++++++++++++++++++--------
 src/context.rs                       | 14 ++----------
 src/parser.rs                        |  7 ++++--
 tests/enums.rs                       |  6 ++---
 6 files changed, 41 insertions(+), 32 deletions(-)

diff --git a/src/abstract_tree/enum_definition.rs b/src/abstract_tree/enum_definition.rs
index a546f4e..5952106 100644
--- a/src/abstract_tree/enum_definition.rs
+++ b/src/abstract_tree/enum_definition.rs
@@ -43,13 +43,13 @@ impl AbstractTree for EnumDefinition {
         Ok(Type::None)
     }
 
-    fn validate(&self, _context: &Context) -> Result<(), ValidationError> {
+    fn validate(&self, context: &Context) -> Result<(), ValidationError> {
+        context.set_enum_definition(self.name.clone(), self.clone())?;
+
         Ok(())
     }
 
-    fn run(self, context: &Context) -> Result<Action, RuntimeError> {
-        context.set_enum_definition(self.name.clone(), self)?;
-
+    fn run(self, _context: &Context) -> Result<Action, RuntimeError> {
         Ok(Action::None)
     }
 }
diff --git a/src/abstract_tree/type.rs b/src/abstract_tree/type.rs
index e1bd1dc..057db14 100644
--- a/src/abstract_tree/type.rs
+++ b/src/abstract_tree/type.rs
@@ -6,7 +6,7 @@ use crate::{
     error::{RuntimeError, TypeConflict, ValidationError},
 };
 
-use super::{AbstractTree, Action};
+use super::{AbstractTree, Action, WithPosition};
 
 #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord)]
 pub enum Type {
@@ -16,7 +16,7 @@ pub enum Type {
     Enum {
         name: Identifier,
         type_arguments: Option<Vec<Type>>,
-        variants: Vec<(Identifier, Option<Type>)>,
+        variants: Vec<(Identifier, Option<WithPosition<Type>>)>,
     },
     Float,
     Function {
diff --git a/src/abstract_tree/value_node.rs b/src/abstract_tree/value_node.rs
index 741b78c..cc7dd02 100644
--- a/src/abstract_tree/value_node.rs
+++ b/src/abstract_tree/value_node.rs
@@ -42,12 +42,7 @@ impl AbstractTree for ValueNode {
         let r#type = match self {
             ValueNode::Boolean(_) => Type::Boolean,
 
-            ValueNode::EnumInstance {
-                name,
-                variant: _,
-                type_arguments,
-                expression: _,
-            } => {
+            ValueNode::EnumInstance { name, .. } => {
                 if let Some(r#type) = context.get_type(name)? {
                     r#type
                 } else {
@@ -130,13 +125,34 @@ impl AbstractTree for ValueNode {
         }
 
         if let ValueNode::EnumInstance {
-            name,
-            variant,
-            type_arguments,
             expression,
+            variant,
+            ..
         } = self
         {
             let r#type = self.expected_type(context)?;
+
+            if let Type::Enum { variants, .. } = r#type {
+                let expected_type_option = variants.iter().find_map(|(identifier, type_option)| {
+                    if identifier == variant {
+                        type_option.clone()
+                    } else {
+                        None
+                    }
+                });
+
+                if let Some(expected_type) = expected_type_option {
+                    let actual_type = expression.node.expected_type(context)?;
+
+                    expected_type.node.check(&actual_type).map_err(|conflict| {
+                        ValidationError::TypeCheck {
+                            conflict,
+                            actual_position: expression.position,
+                            expected_position: expected_type.position,
+                        }
+                    })?
+                }
+            }
         }
 
         Ok(())
diff --git a/src/context.rs b/src/context.rs
index be0cde5..2524851 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -88,26 +88,16 @@ impl Context {
                                     .map(|identifier| Type::Parameter(identifier.clone()))
                                     .collect()
                             });
-                    let variants = enum_definition
-                        .variants()
-                        .into_iter()
-                        .map(|(identifier, type_option)| {
-                            (
-                                identifier.clone(),
-                                type_option.clone().map(|r#type| r#type.node),
-                            )
-                        })
-                        .collect();
 
                     Type::Enum {
                         name: enum_definition.name().clone(),
                         type_arguments,
-                        variants,
+                        variants: enum_definition.variants().clone(),
                     }
                 }
             };
 
-            return Ok(Some(r#type.clone()));
+            return Ok(Some(r#type));
         }
 
         for module in BUILT_IN_MODULES {
diff --git a/src/parser.rs b/src/parser.rs
index 08ea1d4..96dd6ce 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -522,7 +522,7 @@ mod tests {
             parse(
                 &lex("
                 enum FooBar {
-                    Foo,
+                    Foo(int),
                     Bar,    
                 }
                 ")
@@ -534,7 +534,10 @@ mod tests {
                 Identifier::new("FooBar"),
                 None,
                 vec![
-                    (Identifier::new("Foo"), None),
+                    (
+                        Identifier::new("Foo"),
+                        Some(Type::Integer.with_position((55, 58)))
+                    ),
                     (Identifier::new("Bar"), None)
                 ]
             ))
diff --git a/tests/enums.rs b/tests/enums.rs
index 14f4ebe..791fc29 100644
--- a/tests/enums.rs
+++ b/tests/enums.rs
@@ -24,10 +24,10 @@ fn simple_enum_type_check() {
                     actual: Type::String,
                     expected: Type::Integer,
                 },
-                actual_position: (0, 0).into(),
-                expected_position: (0, 0).into()
+                actual_position: (123, 127).into(),
+                expected_position: (47, 50).into()
             },
-            position: (0, 0).into()
+            position: (105, 141).into()
         }])
     )
 }