From 160e4a563639a9a763a278e07faaf443339e2691 Mon Sep 17 00:00:00 2001
From: Josh B <421772+HT154@users.noreply.github.com>
Date: Mon, 13 Jan 2025 16:22:52 -0800
Subject: [PATCH] YamlRenderer: allow all primitive scalar types as map keys
 (#879)

---
 .../org/pkl/core/stdlib/AbstractRenderer.java | 12 ++++++-
 .../core/stdlib/base/YamlRendererNodes.java   | 24 +++++++++++---
 .../core/stdlib/protobuf/RendererNodes.java   |  6 ++--
 .../org/pkl/core/errorMessages.properties     | 10 ++++--
 .../input/api/yamlRenderer8.yml.pkl           | 33 +++++++++++++++++++
 .../output/api/yamlRenderer8.yml              | 26 +++++++++++++++
 6 files changed, 101 insertions(+), 10 deletions(-)
 create mode 100644 pkl-core/src/test/files/LanguageSnippetTests/input/api/yamlRenderer8.yml.pkl
 create mode 100644 pkl-core/src/test/files/LanguageSnippetTests/output/api/yamlRenderer8.yml

diff --git a/pkl-core/src/main/java/org/pkl/core/stdlib/AbstractRenderer.java b/pkl-core/src/main/java/org/pkl/core/stdlib/AbstractRenderer.java
index 10882c26..604868a1 100644
--- a/pkl-core/src/main/java/org/pkl/core/stdlib/AbstractRenderer.java
+++ b/pkl-core/src/main/java/org/pkl/core/stdlib/AbstractRenderer.java
@@ -1,5 +1,5 @@
 /*
- * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ * Copyright © 2024-2025 Apple Inc. and the Pkl project authors. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -415,6 +415,16 @@ public abstract class AbstractRenderer implements VmValueVisitor {
         .build();
   }
 
+  protected void cannotRenderNonScalarKey(Object key) {
+    assert enclosingValue != null;
+    var isMap = enclosingValue instanceof VmMap;
+    throw new VmExceptionBuilder()
+        .evalError(isMap ? "cannotRenderNonScalarMap" : "cannotRenderObjectWithNonScalarKey", name)
+        .withProgramValue(isMap ? "Map" : "Object", enclosingValue)
+        .withProgramValue("Key", key)
+        .build();
+  }
+
   protected void increaseIndent() {
     currIndent.append(indent);
   }
diff --git a/pkl-core/src/main/java/org/pkl/core/stdlib/base/YamlRendererNodes.java b/pkl-core/src/main/java/org/pkl/core/stdlib/base/YamlRendererNodes.java
index 224874c7..cc102172 100644
--- a/pkl-core/src/main/java/org/pkl/core/stdlib/base/YamlRendererNodes.java
+++ b/pkl-core/src/main/java/org/pkl/core/stdlib/base/YamlRendererNodes.java
@@ -1,5 +1,5 @@
 /*
- * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ * Copyright © 2024-2025 Apple Inc. and the Pkl project authors. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -278,8 +278,24 @@ public final class YamlRendererNodes {
         startNewLine();
       }
 
-      if (key instanceof String string) {
-        emitter.emit(string, currIndent, true);
+      if (key instanceof String stringKey) {
+        emitter.emit(stringKey, currIndent, true);
+        builder.append(':');
+        return;
+      } else if (key instanceof Long longKey) {
+        emitter.emit(longKey);
+        builder.append(':');
+        return;
+      } else if (key instanceof Double doubleKey) {
+        emitter.emit(doubleKey);
+        builder.append(':');
+        return;
+      } else if (key instanceof Boolean booleanKey) {
+        emitter.emit(booleanKey);
+        builder.append(':');
+        return;
+      } else if (key instanceof VmNull) {
+        emitter.emitNull();
         builder.append(':');
         return;
       }
@@ -290,7 +306,7 @@ public final class YamlRendererNodes {
         return;
       }
 
-      cannotRenderNonStringKey(key);
+      cannotRenderNonScalarKey(key);
     }
 
     @Override
diff --git a/pkl-core/src/main/java/org/pkl/core/stdlib/protobuf/RendererNodes.java b/pkl-core/src/main/java/org/pkl/core/stdlib/protobuf/RendererNodes.java
index ec24c60a..d7872166 100644
--- a/pkl-core/src/main/java/org/pkl/core/stdlib/protobuf/RendererNodes.java
+++ b/pkl-core/src/main/java/org/pkl/core/stdlib/protobuf/RendererNodes.java
@@ -1,5 +1,5 @@
 /*
- * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ * Copyright © 2024-2025 Apple Inc. and the Pkl project authors. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -310,7 +310,7 @@ public final class RendererNodes {
           isDirective || key instanceof Long || key instanceof Boolean || key instanceof String;
       if (!isValidKey) {
         throw new VmExceptionBuilder()
-            .evalError("cannotRenderNonScalarMapKey")
+            .evalError("cannotRenderProtobufMapKey")
             .withProgramValue("Key", key)
             .build();
       }
@@ -583,7 +583,7 @@ public final class RendererNodes {
             || keyType instanceof StringTypeNode
             || keyType instanceof BooleanTypeNode)) {
           throw new VmExceptionBuilder()
-              .evalError("cannotRenderNonScalarMapKeyType")
+              .evalError("cannotRenderProtobufMapKeyType")
               .withSourceSection(type.getSourceSection())
               .build();
         }
diff --git a/pkl-core/src/main/resources/org/pkl/core/errorMessages.properties b/pkl-core/src/main/resources/org/pkl/core/errorMessages.properties
index 67734035..51c0c2eb 100644
--- a/pkl-core/src/main/resources/org/pkl/core/errorMessages.properties
+++ b/pkl-core/src/main/resources/org/pkl/core/errorMessages.properties
@@ -608,10 +608,16 @@ Cannot render map with non-string key as {0}.
 cannotRenderObjectWithNonStringKey=\
 Cannot render object with non-string key as {0}.
 
-cannotRenderNonScalarMapKey=\
+cannotRenderNonScalarMap=\
+Cannot render map with non-scalar key as {0}.
+
+cannotRenderObjectWithNonScalarKey=\
+Cannot render object with non-scalar key as {0}.
+
+cannotRenderProtobufMapKey=\
 Cannot render map with a non-scalar, floating point or byte array key.
 
-cannotRenderNonScalarMapKeyType=\
+cannotRenderProtobufMapKeyType=\
 Cannot render map with a non-scalar, floating point or byte array key type.
 
 cannotRenderObjectWithElementsAndOtherMembers=\
diff --git a/pkl-core/src/test/files/LanguageSnippetTests/input/api/yamlRenderer8.yml.pkl b/pkl-core/src/test/files/LanguageSnippetTests/input/api/yamlRenderer8.yml.pkl
new file mode 100644
index 00000000..2184bb20
--- /dev/null
+++ b/pkl-core/src/test/files/LanguageSnippetTests/input/api/yamlRenderer8.yml.pkl
@@ -0,0 +1,33 @@
+// Test that non-String scalars can be rendered as YAML map keys.
+mappingInt = new Mapping {
+  [5] = "abc"
+  [10] = "def"
+  [20] = "ghi"
+}
+
+mapInt = mappingInt.toMap()
+
+mappingFloat = new Mapping {
+  [5.0] = "abc"
+  [10.5] = "def"
+  [20.999999990] = "ghi"
+}
+
+mapFloat = mappingFloat.toMap()
+
+mappingBool = new Mapping {
+  [true] = "abc"
+  [false] = "def"
+}
+
+mapBool = mappingBool.toMap()
+
+mappingNull = new Mapping {
+  [null] = "abc"
+}
+
+mapNull = mappingNull.toMap()
+
+output {
+  renderer = new YamlRenderer {}
+}
diff --git a/pkl-core/src/test/files/LanguageSnippetTests/output/api/yamlRenderer8.yml b/pkl-core/src/test/files/LanguageSnippetTests/output/api/yamlRenderer8.yml
new file mode 100644
index 00000000..e4b1b871
--- /dev/null
+++ b/pkl-core/src/test/files/LanguageSnippetTests/output/api/yamlRenderer8.yml
@@ -0,0 +1,26 @@
+mappingInt:
+  5: abc
+  10: def
+  20: ghi
+mapInt:
+  5: abc
+  10: def
+  20: ghi
+mappingFloat:
+  5.0: abc
+  10.5: def
+  20.99999999: ghi
+mapFloat:
+  5.0: abc
+  10.5: def
+  20.99999999: ghi
+mappingBool:
+  true: abc
+  false: def
+mapBool:
+  true: abc
+  false: def
+mappingNull:
+  null: abc
+mapNull:
+  null: abc