4 files changed, 51 insertions, 21 deletions

diff --git a/dhall_core/src/core.rs b/dhall_core/src/core.rs
index 6100981..075ef81 100644
--- a/dhall_core/src/core.rs
+++ b/dhall_core/src/core.rs
@@ -186,10 +186,10 @@ pub enum ExprF<SubExpr, Label, Note, Embed> {
     RecordType(BTreeMap<Label, SubExpr>),
     ///  `{ k1 = v1, k2 = v2 }`
     RecordLit(BTreeMap<Label, SubExpr>),
-    ///  `< k1 : t1, k2 : t2 >`
-    UnionType(BTreeMap<Label, SubExpr>),
-    ///  `< k1 = t1, k2 : t2, k3 : t3 >`
-    UnionLit(Label, SubExpr, BTreeMap<Label, SubExpr>),
+    ///  `< k1 : t1, k2 >`
+    UnionType(BTreeMap<Label, Option<SubExpr>>),
+    ///  `< k1 = t1, k2 : t2, k3 >`
+    UnionLit(Label, SubExpr, BTreeMap<Label, Option<SubExpr>>),
     ///  `merge x y : t`
     Merge(SubExpr, SubExpr, Option<SubExpr>),
     ///  `e.x`
diff --git a/dhall_core/src/parser.rs b/dhall_core/src/parser.rs
index c4ae2e5..a335ccc 100644
--- a/dhall_core/src/parser.rs
+++ b/dhall_core/src/parser.rs
@@ -874,41 +874,41 @@ make_parser! {
 
     rule!(non_empty_union_type_or_literal
           <(Option<(Label, ParsedSubExpr<'a>)>,
-            BTreeMap<Label, ParsedSubExpr<'a>>)>;
+            BTreeMap<Label, Option<ParsedSubExpr<'a>>>)>;
             children!(
         [label(l), union_literal_variant_value((e, entries))] => {
-            (Option::Some((l, rc(e))), entries)
+            (Option::Some((l, e)), entries)
         },
         [label(l), union_type_or_literal_variant_type((e, rest))] => {
             let (x, mut entries) = rest;
-            entries.insert(l, rc(e));
+            entries.insert(l, e);
             (x, entries)
         },
     ));
 
     rule!(union_literal_variant_value
-          <(ParsedExpr<'a>, BTreeMap<Label, ParsedSubExpr<'a>>)>;
+          <(ParsedSubExpr<'a>, BTreeMap<Label, Option<ParsedSubExpr<'a>>>)>;
             children!(
         [expression(e), union_type_entry(entries)..] => {
-            (e, entries.collect())
+            (rc(e), entries.collect())
         },
     ));
 
-    rule!(union_type_entry<(Label, ParsedSubExpr<'a>)>; children!(
-        [label(name), expression(expr)] => (name, rc(expr))
+    rule!(union_type_entry<(Label, Option<ParsedSubExpr<'a>>)>; children!(
+        [label(name), expression(expr)] => (name, Option::Some(rc(expr)))
     ));
 
     // TODO: unary union variants
     rule!(union_type_or_literal_variant_type
-          <(ParsedExpr<'a>,
+          <(Option<ParsedSubExpr<'a>>,
             (Option<(Label, ParsedSubExpr<'a>)>,
-             BTreeMap<Label, ParsedSubExpr<'a>>))>;
+             BTreeMap<Label, Option<ParsedSubExpr<'a>>>))>;
                 children!(
         [expression(e), non_empty_union_type_or_literal(rest)] => {
-            (e, rest)
+            (Option::Some(rc(e)), rest)
         },
         [expression(e)] => {
-            (e, (Option::None, BTreeMap::new()))
+            (Option::Some(rc(e)), (Option::None, BTreeMap::new()))
         },
     ));
 
diff --git a/dhall_core/src/printer.rs b/dhall_core/src/printer.rs
index bb094de..bb3c427 100644
--- a/dhall_core/src/printer.rs
+++ b/dhall_core/src/printer.rs
@@ -92,12 +92,19 @@ impl<SE: Display + Clone, N, E: Display> Display for ExprF<SE, Label, N, E> {
                 write!(f, "{} = {}", k, v)
             })?,
             UnionType(a) => fmt_list("< ", " | ", " >", a, f, |(k, v), f| {
-                write!(f, "{} : {}", k, v)
+                write!(f, "{} : ", k)?;
+                if let Some(v) = v {
+                    v.fmt(f)?
+                }
+                Ok(())
             })?,
             UnionLit(a, b, c) => {
                 write!(f, "< {} = {}", a, b)?;
                 for (k, v) in c {
-                    write!(f, " | {} : {}", k, v)?;
+                    write!(f, " | {}", k)?;
+                    if let Some(v) = v {
+                        write!(f, ": {}", v)?;
+                    }
                 }
                 f.write_str(" >")?
             }
diff --git a/dhall_core/src/visitor.rs b/dhall_core/src/visitor.rs
index 3b06f8b..1ea3fb1 100644
--- a/dhall_core/src/visitor.rs
+++ b/dhall_core/src/visitor.rs
@@ -106,6 +106,27 @@ where
                 .map(|(k, x)| Ok((v.visit_label(k)?, v.visit_subexpr(x)?)))
                 .collect()
         }
+        fn btoptmap<'a, V, SE, L, N, E, SE2, L2, N2, E2>(
+            x: &'a BTreeMap<L, Option<SE>>,
+            mut v: V,
+        ) -> Result<BTreeMap<L2, Option<SE2>>, V::Error>
+        where
+            L: Ord,
+            L2: Ord,
+            V: ExprFFallibleVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>,
+        {
+            x.iter()
+                .map(|(k, x)| {
+                    Ok((
+                        v.visit_label(k)?,
+                        match x {
+                            Some(x) => Some(v.visit_subexpr(x)?),
+                            None => None,
+                        },
+                    ))
+                })
+                .collect()
+        }
 
         let mut v = self;
         use crate::ExprF::*;
@@ -156,10 +177,12 @@ where
             NEOptionalLit(e) => NEOptionalLit(v.visit_subexpr(e)?),
             RecordType(kts) => RecordType(btmap(kts, v)?),
             RecordLit(kvs) => RecordLit(btmap(kvs, v)?),
-            UnionType(kts) => UnionType(btmap(kts, v)?),
-            UnionLit(k, x, kvs) => {
-                UnionLit(v.visit_label(k)?, v.visit_subexpr(x)?, btmap(kvs, v)?)
-            }
+            UnionType(kts) => UnionType(btoptmap(kts, v)?),
+            UnionLit(k, x, kvs) => UnionLit(
+                v.visit_label(k)?,
+                v.visit_subexpr(x)?,
+                btoptmap(kvs, v)?,
+            ),
             Merge(x, y, t) => Merge(
                 v.visit_subexpr(x)?,
                 v.visit_subexpr(y)?,