HIR - Union handling

12 files changed, 176 insertions, 16 deletions

diff --git a/src/hir/dump.cpp b/src/hir/dump.cpp
index 8139bcd5..853d62f1 100644
--- a/src/hir/dump.cpp
+++ b/src/hir/dump.cpp
@@ -610,6 +610,12 @@ namespace {
             m_os << indent() << "}";
             dec_indent();
         }
+        void visit(::HIR::ExprNode_UnionLiteral& node) override
+        {
+            m_os << node.m_path << " { " << node.m_variant_name << ": ";
+            this->visit_node_ptr( node.m_value );
+            m_os << " }";
+        }
         void visit(::HIR::ExprNode_Tuple& node) override
         {
             m_os << "(";
diff --git a/src/hir/expr.cpp b/src/hir/expr.cpp
index e11002d9..02e7ed98 100644
--- a/src/hir/expr.cpp
+++ b/src/hir/expr.cpp
@@ -141,6 +141,10 @@ DEF_VISIT(ExprNode_StructLiteral, node,
     for(auto& val : node.m_values)
         visit_node_ptr(val.second);
 )
+DEF_VISIT(ExprNode_UnionLiteral, node,
+    visit_generic_path(::HIR::Visitor::PathContext::TYPE, node.m_path);
+    visit_node_ptr(node.m_value);
+)
 DEF_VISIT(ExprNode_Tuple, node,
     for(auto& val : node.m_vals)
         visit_node_ptr(val);
diff --git a/src/hir/expr.hpp b/src/hir/expr.hpp
index cd4a21ab..a5a1346a 100644
--- a/src/hir/expr.hpp
+++ b/src/hir/expr.hpp
@@ -690,6 +690,25 @@ struct ExprNode_StructLiteral:
     
     NODE_METHODS();
 };
+struct ExprNode_UnionLiteral:
+    public ExprNode
+{
+    ::HIR::GenericPath  m_path;
+    ::std::string   m_variant_name;
+    ::HIR::ExprNodeP    m_value;
+    
+    unsigned int    m_variant_index = ~0;
+    
+    ExprNode_UnionLiteral(Span sp, ::HIR::GenericPath path, ::std::string name, ::HIR::ExprNodeP value):
+        ExprNode( mv$(sp) ),
+        m_path( mv$(path) ),
+        m_variant_name( mv$(name) ),
+        m_value( mv$(value) )
+    {
+    }
+    
+    NODE_METHODS();
+};
 struct ExprNode_Tuple:
     public ExprNode
 {
@@ -803,6 +822,7 @@ public:
     NV(ExprNode_Variable);
     
     NV(ExprNode_StructLiteral);
+    NV(ExprNode_UnionLiteral);
     NV(ExprNode_Tuple);
     NV(ExprNode_ArrayList);
     NV(ExprNode_ArraySized);
@@ -847,6 +867,7 @@ public:
     NV(ExprNode_Variable);
     
     NV(ExprNode_StructLiteral);
+    NV(ExprNode_UnionLiteral);
     NV(ExprNode_Tuple);
     NV(ExprNode_ArrayList);
     NV(ExprNode_ArraySized);
diff --git a/src/hir/from_ast_expr.cpp b/src/hir/from_ast_expr.cpp
index 21024dd7..4c10e91a 100644
--- a/src/hir/from_ast_expr.cpp
+++ b/src/hir/from_ast_expr.cpp
@@ -537,15 +537,31 @@ struct LowerHIR_ExprNode_Visitor:
             ) );
     }
     virtual void visit(::AST::ExprNode_StructLiteral& v) override {
-        ::HIR::ExprNode_StructLiteral::t_values values;
-        for(const auto& val : v.m_values)
-            values.push_back( ::std::make_pair(val.first, LowerHIR_ExprNode_Inner(*val.second)) );
-        m_rv.reset( new ::HIR::ExprNode_StructLiteral( v.span(),
-            LowerHIR_GenericPath(v.get_pos(), v.m_path),
-            ! v.m_path.binding().is_EnumVar(),
-            LowerHIR_ExprNode_Inner_Opt(v.m_base_value.get()),
-            mv$(values)
-            ) );
+        if( v.m_path.binding().is_Union() )
+        {
+            if( v.m_values.size() != 1 )
+                ERROR(v.span(), E0000, "Union constructors can only specify a single field");
+            if( v.m_base_value )
+                ERROR(v.span(), E0000, "Union constructors can't take a base value");
+            
+            m_rv.reset( new ::HIR::ExprNode_UnionLiteral( v.span(),
+                LowerHIR_GenericPath(v.get_pos(), v.m_path),
+                v.m_values[0].first,
+                LowerHIR_ExprNode_Inner(*v.m_values[0].second)
+                ) );
+        }
+        else
+        {
+            ::HIR::ExprNode_StructLiteral::t_values values;
+            for(const auto& val : v.m_values)
+                values.push_back( ::std::make_pair(val.first, LowerHIR_ExprNode_Inner(*val.second)) );
+            m_rv.reset( new ::HIR::ExprNode_StructLiteral( v.span(),
+                LowerHIR_GenericPath(v.get_pos(), v.m_path),
+                ! v.m_path.binding().is_EnumVar(),
+                LowerHIR_ExprNode_Inner_Opt(v.m_base_value.get()),
+                mv$(values)
+                ) );
+        }
     }
     virtual void visit(::AST::ExprNode_Array& v) override {
         if( v.m_size )
diff --git a/src/hir/hir.cpp b/src/hir/hir.cpp
index 0c42ba05..512f6c5d 100644
--- a/src/hir/hir.cpp
+++ b/src/hir/hir.cpp
@@ -522,6 +522,16 @@ const ::HIR::Struct& ::HIR::Crate::get_struct_by_path(const Span& sp, const ::HI
         BUG(sp, "Struct path " << path << " didn't point to a struct");
     }
 }
+const ::HIR::Union& ::HIR::Crate::get_union_by_path(const Span& sp, const ::HIR::SimplePath& path) const
+{
+    const auto& ti = this->get_typeitem_by_path(sp, path);
+    TU_IFLET(::HIR::TypeItem, ti, Union, e,
+        return e;
+    )
+    else {
+        BUG(sp, "Path " << path << " didn't point to a union");
+    }
+}
 const ::HIR::Enum& ::HIR::Crate::get_enum_by_path(const Span& sp, const ::HIR::SimplePath& path) const
 {
     const auto& ti = this->get_typeitem_by_path(sp, path);
diff --git a/src/hir/hir.hpp b/src/hir/hir.hpp
index 2e0e6592..d5f04fd3 100644
--- a/src/hir/hir.hpp
+++ b/src/hir/hir.hpp
@@ -389,6 +389,7 @@ public:
     const ::HIR::TypeItem& get_typeitem_by_path(const Span& sp, const ::HIR::SimplePath& path, bool ignore_crate_name=false) const;
     const ::HIR::Trait& get_trait_by_path(const Span& sp, const ::HIR::SimplePath& path) const;
     const ::HIR::Struct& get_struct_by_path(const Span& sp, const ::HIR::SimplePath& path) const;
+    const ::HIR::Union& get_union_by_path(const Span& sp, const ::HIR::SimplePath& path) const;
     const ::HIR::Enum& get_enum_by_path(const Span& sp, const ::HIR::SimplePath& path) const;
     const ::HIR::Module& get_mod_by_path(const Span& sp, const ::HIR::SimplePath& path) const;
     
diff --git a/src/hir_conv/constant_evaluation.cpp b/src/hir_conv/constant_evaluation.cpp
index dfb008e3..8d93ced1 100644
--- a/src/hir_conv/constant_evaluation.cpp
+++ b/src/hir_conv/constant_evaluation.cpp
@@ -818,6 +818,10 @@ namespace {
                     TODO(node.span(), "Handle Enum _UnitVariant - " << node.m_path);
                 }
             }
+            void visit(::HIR::ExprNode_UnionLiteral& node) override {
+                TRACE_FUNCTION_FR("_UnionLiteral - " << node.m_path, m_rv);
+                TODO(node.span(), "_UnionLiteral");
+            }
             void visit(::HIR::ExprNode_Tuple& node) override {
                 ::std::vector< ::HIR::Literal>  vals;
                 ::std::vector< ::HIR::TypeRef>  tys;
diff --git a/src/hir_expand/annotate_value_usage.cpp b/src/hir_expand/annotate_value_usage.cpp
index 7f2462d9..5ddfd39b 100644
--- a/src/hir_expand/annotate_value_usage.cpp
+++ b/src/hir_expand/annotate_value_usage.cpp
@@ -355,6 +355,11 @@ namespace {
                 this->visit_node_ptr(fld_val.second);
             }
         }
+        void visit(::HIR::ExprNode_UnionLiteral& node) override
+        {
+            auto _ = push_usage( ::HIR::ValueUsage::Move );
+            this->visit_node_ptr(node.m_value);
+        }
         void visit(::HIR::ExprNode_Tuple& node) override
         {
             auto _ = push_usage( ::HIR::ValueUsage::Move );
diff --git a/src/hir_typeck/expr_check.cpp b/src/hir_typeck/expr_check.cpp
index 6dde6e84..7af82bb1 100644
--- a/src/hir_typeck/expr_check.cpp
+++ b/src/hir_typeck/expr_check.cpp
@@ -522,12 +522,17 @@ namespace {
                 node.m_base_value->visit( *this );
             }
         }
+        void visit(::HIR::ExprNode_UnionLiteral& node) override
+        {
+            // TODO: Check.
+            node.m_value->visit(*this);
+        }
         void visit(::HIR::ExprNode_UnitVariant& node) override
         {
             TRACE_FUNCTION_F(&node << " " << node.m_path << " [" << (node.m_is_struct ? "struct" : "enum") << "]");
             const auto& sp = node.span();
             const auto& ty = node.m_res_type;
-            ASSERT_BUG(sp, ty.m_data.is_Path(), "Result type of _StructLiteral isn't Path");
+            ASSERT_BUG(sp, ty.m_data.is_Path(), "Result type of _UnitVariant isn't Path");
             
             TU_MATCH(::HIR::TypeRef::TypePathBinding, (ty.m_data.as_Path().binding), (e),
             (Unbound, ),
@@ -870,10 +875,18 @@ namespace {
             {
                 ASSERT_BUG(sp, str_ty.m_data.is_Path(), "Value type of _Field isn't Path - " << str_ty);
                 const auto& ty_e = str_ty.m_data.as_Path();
-                ASSERT_BUG(sp, ty_e.binding.is_Struct(), "Value type of _Field isn't a Struct - " << str_ty);
-                //const auto& str = *ty_e.binding.as_Struct();
-                
-                // TODO: Triple-check result, but that probably isn't needed
+                if( ty_e.binding.is_Struct() )
+                {
+                    //const auto& str = *ty_e.binding.as_Struct();
+                    // TODO: Triple-check result, but that probably isn't needed
+                }
+                else if( ty_e.binding.is_Union() )
+                {
+                }
+                else
+                {
+                    ASSERT_BUG(sp, ty_e.binding.is_Struct() || ty_e.binding.is_Union(), "Value type of _Field isn't a Struct or Union - " << str_ty);
+                }
             }
             
             node.m_value->visit( *this );
diff --git a/src/hir_typeck/expr_cs.cpp b/src/hir_typeck/expr_cs.cpp
index 5836cbef..97b5ae72 100644
--- a/src/hir_typeck/expr_cs.cpp
+++ b/src/hir_typeck/expr_cs.cpp
@@ -1208,6 +1208,56 @@ namespace {
                 node.m_base_value->visit( *this );
             }
         }
+        void visit(::HIR::ExprNode_UnionLiteral& node) override
+        {
+            const Span& sp = node.span();
+            TRACE_FUNCTION_F(&node << " " << node.m_path << "{ " << node.m_variant_name << ": ... }");
+            this->context.add_ivars( node.m_value->m_res_type );
+            
+            const auto& unm = this->context.m_crate.get_union_by_path(sp, node.m_path.m_path);
+            fix_param_count(sp, this->context, ::HIR::TypeRef(), false, node.m_path, unm.m_params, node.m_path.m_params);
+            const auto ty = ::HIR::TypeRef::new_path( node.m_path.clone(), ::HIR::TypeRef::TypePathBinding::make_Union(&unm) );
+            
+            this->context.equate_types(node.span(), node.m_res_type, ty);
+            
+            const auto& ty_params = node.m_path.m_params.m_types;
+            auto monomorph_cb = [&](const auto& gt)->const auto& {
+                const auto& ge = gt.m_data.as_Generic();
+                if( ge.binding == 0xFFFF ) {
+                    return ty;
+                }
+                else if( ge.binding < 256 ) {
+                    if( ge.binding >= ty_params.size() ) {
+                        BUG(node.span(), "Type parameter index out of range (#" << ge.binding << " " << ge.name << ")");
+                    }
+                    return ty_params[ge.binding];
+                }
+                else {
+                    BUG(node.span(), "Method-level parameter on struct (#" << ge.binding << " " << ge.name << ")");
+                }
+                };
+            
+            // Convert bounds on the type into rules
+            apply_bounds_as_rules(context, node.span(), unm.m_params, monomorph_cb);
+
+            auto it = ::std::find_if(unm.m_variants.begin(), unm.m_variants.end(), [&](const auto& v)->bool{ return v.first == node.m_variant_name; });
+            assert(it != unm.m_variants.end());
+            const auto& des_ty_r = it->second.ent;
+            ::HIR::TypeRef  des_ty_cache;
+            const auto* des_ty = &des_ty_r;
+            if( monomorphise_type_needed(des_ty_r) ) {
+                if( des_ty_cache == ::HIR::TypeRef() ) {
+                    des_ty_cache = monomorphise_type_with(node.span(), des_ty_r, monomorph_cb);
+                }
+                else {
+                    // TODO: Is it an error when it's already populated?
+                }
+                des_ty = &des_ty_cache;
+            }
+            this->equate_types_inner_coerce(node.span(), *des_ty,  node.m_value);
+            
+            node.m_value->visit(*this);
+        }
         void visit(::HIR::ExprNode_UnitVariant& node) override
         {
             TRACE_FUNCTION_F(&node << " " << node.m_path << " [" << (node.m_is_struct ? "struct" : "enum") << "]");
@@ -2732,6 +2782,9 @@ namespace {
         void visit(::HIR::ExprNode_StructLiteral& node) override {
             no_revisit(node);
         }
+        void visit(::HIR::ExprNode_UnionLiteral& node) override {
+            no_revisit(node);
+        }
         void visit(::HIR::ExprNode_Tuple& node) override {
             no_revisit(node);
         }
diff --git a/src/hir_typeck/helpers.cpp b/src/hir_typeck/helpers.cpp
index 88a4624f..c7e18fe4 100644
--- a/src/hir_typeck/helpers.cpp
+++ b/src/hir_typeck/helpers.cpp
@@ -3569,7 +3569,30 @@ bool TraitResolution::find_field(const Span& sp, const ::HIR::TypeRef& ty, const
             // No fields on enums either
             ),
         (Union,
-            TODO(sp, "Field search on union");
+            const auto& unm = *be;
+            const auto& params = e.path.m_data.as_Generic().m_params;
+            auto monomorph = [&](const auto& gt)->const auto& {
+                const auto& ge = gt.m_data.as_Generic();
+                if( ge.binding == 0xFFFF )
+                    TODO(sp, "Monomorphise union field types (Self) - " << gt);
+                else if( ge.binding < 256 ) {
+                    assert(ge.binding < params.m_types.size());
+                    return params.m_types[ge.binding];
+                }
+                else {
+                    BUG(sp, "function-level param encountered in union field");
+                }
+                return gt;
+                };
+            
+            for( const auto& fld : unm.m_variants )
+            {
+                // TODO: Privacy
+                if( fld.first == name ) {
+                    field_ty = monomorphise_type_with(sp, fld.second.ent, monomorph);
+                    return true;
+                }
+            }
             )
         )
     )
diff --git a/src/mir/from_hir.cpp b/src/mir/from_hir.cpp
index 9f8a91e3..d635fc82 100644
--- a/src/mir/from_hir.cpp
+++ b/src/mir/from_hir.cpp
@@ -1668,7 +1668,7 @@ namespace {
                 fields_ptr = &it->second.as_Struct();
                 ),
             (Union,
-                TODO(node.span(), "_StructLiteral Union");
+                BUG(node.span(), "_StructLiteral Union");
                 ),
             (Struct,
                 fields_ptr = &e->m_data.as_Named();
@@ -1709,6 +1709,10 @@ namespace {
                 mv$(values)
                 }) );
         }
+        void visit(::HIR::ExprNode_UnionLiteral& node) override
+        {
+            TODO(node.span(), "_UnionLiteral");
+        }
         
         void visit(::HIR::ExprNode_Tuple& node) override
         {