From 735c443b3750140a41690b2173e5edc05fa7416e Mon Sep 17 00:00:00 2001
From: Mattias Wadman <mattias.wadman@gmail.com>
Date: Sat, 28 May 2022 20:25:20 +0200
Subject: [PATCH] interp: Improve type normalization and use it for toyaml and
 totoml

Use smallest int type for int64, uint6 and *big.Int

Fixes integer serialization for yaml and toml for small integers, othweise
they will end up as strings.
---
 internal/gojqextra/totype.go             | 19 +++++++++++++++++++
 pkg/interp/encoding.go                   |  4 ++--
 pkg/interp/testdata/encoding/toml.fqtest |  1 +
 pkg/interp/testdata/encoding/yaml.fqtest |  1 +
 4 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/internal/gojqextra/totype.go b/internal/gojqextra/totype.go
index 2532b4aa..174cb0f4 100644
--- a/internal/gojqextra/totype.go
+++ b/internal/gojqextra/totype.go
@@ -137,14 +137,33 @@ func ToGoJQValue(v any) (any, bool) {
 	case int:
 		return vv, true
 	case int64:
+		if vv >= math.MinInt && vv <= math.MaxInt {
+			return int(vv), true
+		}
 		return big.NewInt(vv), true
 	case uint64:
+		if vv <= math.MaxInt {
+			return int(vv), true
+		}
 		return new(big.Int).SetUint64(vv), true
 	case float32:
 		return float64(vv), true
 	case float64:
 		return vv, true
 	case *big.Int:
+		if vv.IsInt64() {
+			vv := vv.Int64()
+			if vv >= math.MinInt && vv <= math.MaxInt {
+				return int(vv), true
+			}
+			return vv, true
+		} else if vv.IsUint64() {
+			vv := vv.Uint64()
+			if vv <= math.MaxInt {
+				return int(vv), true
+			}
+			return vv, true
+		}
 		return vv, true
 	case string:
 		return vv, true
diff --git a/pkg/interp/encoding.go b/pkg/interp/encoding.go
index 658670e3..48046ee7 100644
--- a/pkg/interp/encoding.go
+++ b/pkg/interp/encoding.go
@@ -717,7 +717,7 @@ func init() {
 	})
 
 	addFunc("_toyaml", func(c any) any {
-		b, err := yaml.Marshal(c)
+		b, err := yaml.Marshal(norm(c))
 		if err != nil {
 			return err
 		}
@@ -734,7 +734,7 @@ func init() {
 
 	addFunc("_totoml", func(c map[string]any) any {
 		b := &bytes.Buffer{}
-		if err := toml.NewEncoder(b).Encode(c); err != nil {
+		if err := toml.NewEncoder(b).Encode(norm(c)); err != nil {
 			return err
 		}
 		return b.String()
diff --git a/pkg/interp/testdata/encoding/toml.fqtest b/pkg/interp/testdata/encoding/toml.fqtest
index bacfaf65..d16fec83 100644
--- a/pkg/interp/testdata/encoding/toml.fqtest
+++ b/pkg/interp/testdata/encoding/toml.fqtest
@@ -1,4 +1,5 @@
 # toml does not support null in arrays
+# TODO: add uint64 norm test
 $ fq -rRs 'fromjson[] | (walk(if type == "array" then map(select(. != null)) end) | try (totoml | ., fromtoml) catch .), "----"' variants.json
 
 {}
diff --git a/pkg/interp/testdata/encoding/yaml.fqtest b/pkg/interp/testdata/encoding/yaml.fqtest
index a99e2baa..37da05ca 100644
--- a/pkg/interp/testdata/encoding/yaml.fqtest
+++ b/pkg/interp/testdata/encoding/yaml.fqtest
@@ -1,3 +1,4 @@
+# TODO: add uint64 norm test
 $ fq -rRs 'fromjson[] | (try (toyaml | ., fromyaml) catch .), "----"' variants.json
 null