New implementations for with

modules/core/src/main/java/org/dhallj/core/Expr.java

@@ -467,35 +467,6 @@ public static final Expr desugarComplete(Expr lhs, Expr rhs) {
           Expr.makeFieldAccess(lhs, "Type"));
     }
 
-    /** Desugar {@code with}. */
-    public static final Expr desugarWith(Expr base, String[] path, Expr value) {
-      if (path.length == 1) {
-        return Expr.makeOperatorApplication(
-            Operator.PREFER, base, Expr.makeRecordLiteral(path[0], value));
-      } else {
-        Expr[] accessors = new Expr[path.length - 1];
-        accessors[0] = Expr.makeFieldAccess(Constants.UNDERSCORE, path[0]);
-
-        for (int i = 1; i < path.length - 1; i += 1) {
-          accessors[i] = Expr.makeFieldAccess(accessors[i - 1], path[i]);
-        }
-
-        Expr incremented = value.increment("_");
-        Expr current = Expr.makeRecordLiteral(path[path.length - 1], incremented);
-
-        for (int i = path.length - 2; i >= 0; i -= 1) {
-          current =
-              Expr.makeRecordLiteral(
-                  path[i], Expr.makeOperatorApplication(Operator.PREFER, accessors[i], current));
-        }
-
-        return Expr.makeLet(
-            "_",
-            base,
-            Expr.makeOperatorApplication(Operator.PREFER, Constants.UNDERSCORE, current));
-      }
-    }
-
     /**
      * If the expression is a lambda, apply it to the given argument.
      *

modules/core/src/main/java/org/dhallj/core/normalization/BetaNormalize.java

@@ -142,7 +142,7 @@ public Expr onToMap(Expr base, Expr type) {
   }
 
   public Expr onWith(Expr base, String[] path, Expr value) {
-    return Expr.Util.desugarWith(base, path, value).accept(this);
+    return BetaNormalizeWith.apply(base, path, value);
   }
 
   public Expr onMissingImport(Expr.ImportMode mode, byte[] hash) {

modules/core/src/main/java/org/dhallj/core/normalization/BetaNormalizeWith.java

@@ -0,0 +1,78 @@
+package org.dhallj.core.normalization;
+
+import java.util.AbstractMap.SimpleImmutableEntry;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map.Entry;
+import org.dhallj.core.Expr;
+
+final class BetaNormalizeWith {
+  static final Expr apply(Expr base, String[] path, Expr value) {
+    List<Entry<String, Expr>> baseAsRecordLiteral = Expr.Util.asRecordLiteral(base);
+
+    if (baseAsRecordLiteral != null) {
+      List<Entry<String, Expr>> result = new ArrayList<>();
+      boolean found = false;
+
+      for (Entry<String, Expr> field : baseAsRecordLiteral) {
+        String key = field.getKey();
+
+        if (key.equals(path[0])) {
+          found = true;
+
+          Expr newValue;
+
+          if (path.length == 1) {
+            newValue = value;
+          } else {
+            String[] remainingPath = new String[path.length - 1];
+            System.arraycopy(path, 1, remainingPath, 0, path.length - 1);
+
+            newValue =
+                Expr.makeWith(field.getValue(), remainingPath, value)
+                    .accept(BetaNormalize.instance);
+          }
+
+          result.add(new SimpleImmutableEntry<>(key, newValue));
+        } else {
+          result.add(field);
+        }
+      }
+
+      if (!found) {
+        Expr newValue;
+
+        if (path.length == 1) {
+          newValue = value;
+        } else {
+          String[] remainingPath = new String[path.length - 1];
+          System.arraycopy(path, 1, remainingPath, 0, path.length - 1);
+
+          newValue =
+              Expr.makeWith(Expr.Constants.EMPTY_RECORD_LITERAL, remainingPath, value)
+                  .accept(BetaNormalize.instance);
+        }
+
+        result.add(new SimpleImmutableEntry<>(path[0], newValue));
+      }
+
+      Entry<String, Expr>[] resultArray = result.toArray(new Entry[result.size()]);
+
+      Arrays.sort(resultArray, entryComparator);
+
+      return Expr.makeRecordLiteral(resultArray);
+    } else {
+      return Expr.makeWith(base, path, value);
+    }
+  }
+
+  /** Java 8 introduce {@code comparingByKey}, but we can roll our own pretty easily. */
+  private static final Comparator<Entry<String, Expr>> entryComparator =
+      new Comparator<Entry<String, Expr>>() {
+        public int compare(Entry<String, Expr> a, Entry<String, Expr> b) {
+          return a.getKey().compareTo(b.getKey());
+        }
+      };
+}

modules/core/src/main/java/org/dhallj/core/typechecking/TypeCheck.java

@@ -676,7 +676,59 @@ public final Expr onToMap(Expr base, Expr type) {
 
   @Override
   public final Expr onWith(Expr base, String[] path, Expr value) {
-    return Expr.Util.desugarWith(base, path, value).accept(this);
+    List<Entry<String, Expr>> baseAsRecordLiteral = Expr.Util.asRecordLiteral(base);
+
+    if (baseAsRecordLiteral == null) {
+      throw TypeCheckFailure.makeWithTypeError(base.accept(this));
+    } else {
+      List<Entry<String, Expr>> result = new ArrayList<>();
+      boolean found = false;
+
+      for (Entry<String, Expr> field : baseAsRecordLiteral) {
+        String key = field.getKey();
+
+        if (key.equals(path[0])) {
+          found = true;
+
+          Expr newValue;
+
+          if (path.length == 1) {
+            newValue = value.accept(this);
+          } else {
+            String[] remainingPath = new String[path.length - 1];
+            System.arraycopy(path, 1, remainingPath, 0, path.length - 1);
+
+            newValue = Expr.makeWith(field.getValue(), remainingPath, value).accept(this);
+          }
+
+          result.add(new SimpleImmutableEntry<>(key, newValue));
+        } else {
+          result.add(new SimpleImmutableEntry<>(key, field.getValue().accept(this)));
+        }
+      }
+
+      if (!found) {
+        Expr newValue;
+
+        if (path.length == 1) {
+          newValue = value.accept(this);
+        } else {
+          String[] remainingPath = new String[path.length - 1];
+          System.arraycopy(path, 1, remainingPath, 0, path.length - 1);
+
+          newValue =
+              Expr.makeWith(Expr.Constants.EMPTY_RECORD_LITERAL, remainingPath, value).accept(this);
+        }
+
+        result.add(new SimpleImmutableEntry<>(path[0], newValue));
+      }
+
+      Entry<String, Expr>[] resultArray = result.toArray(new Entry[result.size()]);
+
+      Arrays.sort(resultArray, entryComparator);
+
+      return Expr.makeRecordType(resultArray);
+    }
   }
 
   @Override

modules/core/src/main/java/org/dhallj/core/typechecking/TypeCheckFailure.java

@@ -207,4 +207,8 @@ static final TypeCheckFailure makeToMapMissingAnnotationError() {
   static final TypeCheckFailure makeToMapInvalidAnnotationError(Expr type) {
     return new TypeCheckFailure("An empty toMap was annotated with an invalid type");
   }
+
+  static final TypeCheckFailure makeWithTypeError(Expr type) {
+    return new TypeCheckFailure("with only works on records");
+  }
 }