From d37c618ead16772fc2efbcd630eff9556ce50bc8 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Wed, 19 Jun 2024 02:32:17 -0400
Subject: [PATCH] Continue implementing type inference

---
 dust-lang/src/abstract_tree/assignment.rs    | 40 ++++++++++++++++++--
 dust-lang/src/abstract_tree/function_call.rs |  6 ++-
 dust-lang/src/value.rs                       |  4 ++
 examples/type_inference.ds                   |  2 +
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/dust-lang/src/abstract_tree/assignment.rs b/dust-lang/src/abstract_tree/assignment.rs
index 55eb79e..33c01b2 100644
--- a/dust-lang/src/abstract_tree/assignment.rs
+++ b/dust-lang/src/abstract_tree/assignment.rs
@@ -8,7 +8,8 @@ use crate::{
 };
 
 use super::{
-    AbstractNode, Evaluation, ExpectedType, Statement, Type, TypeConstructor, WithPosition,
+    AbstractNode, Evaluation, ExpectedType, Expression, Statement, Type, TypeConstructor,
+    WithPosition,
 };
 
 #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Serialize, Deserialize)]
@@ -52,7 +53,38 @@ impl AbstractNode for Assignment {
             ));
         }
 
-        if let Some(constructor) = &self.constructor {
+        if let (Some(constructor), Statement::Expression(Expression::FunctionCall(function_call))) =
+            (&self.constructor, self.statement.as_ref())
+        {
+            let declared_type = constructor.clone().construct(context)?;
+            let function_type = function_call.node.function().expected_type(context)?;
+
+            if let Type::Function {
+                return_type,
+                type_parameters: Some(type_parameters),
+                ..
+            } = function_type
+            {
+                if let Type::Generic {
+                    identifier,
+                    concrete_type,
+                } = *return_type
+                {
+                    let returned_parameter = type_parameters
+                        .into_iter()
+                        .find(|parameter| parameter == &identifier);
+
+                    if let Some(parameter) = returned_parameter {
+                        context.set_type(parameter, declared_type)?;
+                    }
+                }
+            } else {
+                return Err(ValidationError::ExpectedFunction {
+                    actual: function_type,
+                    position: function_call.position,
+                });
+            }
+        } else if let Some(constructor) = &self.constructor {
             let r#type = constructor.clone().construct(&context)?;
 
             r#type
@@ -78,8 +110,8 @@ impl AbstractNode for Assignment {
         context: &mut Context,
         manage_memory: bool,
     ) -> Result<Evaluation, RuntimeError> {
-        let action = self.statement.evaluate(context, manage_memory)?;
-        let right = match action {
+        let evaluation = self.statement.evaluate(context, manage_memory)?;
+        let right = match evaluation {
             Evaluation::Return(value) => value,
             r#break => return Ok(r#break),
         };
diff --git a/dust-lang/src/abstract_tree/function_call.rs b/dust-lang/src/abstract_tree/function_call.rs
index 18ce91b..1300923 100644
--- a/dust-lang/src/abstract_tree/function_call.rs
+++ b/dust-lang/src/abstract_tree/function_call.rs
@@ -27,6 +27,10 @@ impl FunctionCall {
             value_arguments,
         }
     }
+
+    pub fn function(&self) -> &Box<Expression> {
+        &self.function
+    }
 }
 
 impl AbstractNode for FunctionCall {
@@ -148,7 +152,7 @@ impl ExpectedType for FunctionCall {
                     for (constructor, identifier) in
                         type_arguments.into_iter().zip(type_parameters.into_iter())
                     {
-                        if &return_identifier == identifier {
+                        if identifier == &return_identifier {
                             let concrete_type = constructor.clone().construct(&context)?;
 
                             return Ok(Type::Generic {
diff --git a/dust-lang/src/value.rs b/dust-lang/src/value.rs
index 55703ef..ceb5ec0 100644
--- a/dust-lang/src/value.rs
+++ b/dust-lang/src/value.rs
@@ -341,6 +341,10 @@ impl Function {
         &self.value_parameters
     }
 
+    pub fn body(&self) -> &Block {
+        &self.body
+    }
+
     pub fn call(
         self,
         arguments: Vec<Value>,
diff --git a/examples/type_inference.ds b/examples/type_inference.ds
index 60d7a6b..1cce8b1 100644
--- a/examples/type_inference.ds
+++ b/examples/type_inference.ds
@@ -1,3 +1,5 @@
 foo = fn |T| (x: T) -> T { x }
 bar: str = foo::(str)::("hi")
 baz: str = foo("hi")
+x = json.parse::(int)::("1")
+x: int = json.parse("1")