From 0d4a0f37c92e7f0308568fd69b9f8bdc982811dc Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Tue, 22 Oct 2024 14:08:12 +0200
Subject: [PATCH] Fix allownil on maps elements

Similar to #374 map `[]byte` elements would also inherit allownil unintentionally.
---
 gen/decode.go    | 1 +
 gen/elem.go      | 4 ++++
 gen/encode.go    | 1 +
 gen/marshal.go   | 1 +
 gen/spec.go      | 7 +------
 gen/unmarshal.go | 1 +
 6 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/gen/decode.go b/gen/decode.go
index daca62be..eb060e26 100644
--- a/gen/decode.go
+++ b/gen/decode.go
@@ -257,6 +257,7 @@ func (d *decodeGen) gMap(m *Map) {
 	d.p.declare(m.Validx, m.Value.TypeName())
 	d.assignAndCheck(m.Keyidx, stringTyp)
 	d.ctx.PushVar(m.Keyidx)
+	m.Value.SetIsAllowNil(false)
 	next(d, m.Value)
 	d.p.mapAssign(m)
 	d.ctx.Pop()
diff --git a/gen/elem.go b/gen/elem.go
index 1455170f..153f9061 100644
--- a/gen/elem.go
+++ b/gen/elem.go
@@ -146,6 +146,7 @@ func (c *common) Varname() string     { return c.vname }
 func (c *common) Alias(typ string)    { c.alias = typ }
 func (c *common) hidden()             {}
 func (c *common) AllowNil() bool      { return false }
+func (c *common) SetIsAllowNil(bool)  {}
 func (c *common) AlwaysPtr(set *bool) bool {
 	if c != nil && set != nil {
 		c.ptrRcv = *set
@@ -202,6 +203,9 @@ type Elem interface {
 	// This is true for slices and maps.
 	AllowNil() bool
 
+	// SetIsAllowNil will set the allownil value, if the type supports it.
+	SetIsAllowNil(bool)
+
 	// AlwaysPtr will return true if receiver should always be a pointer.
 	AlwaysPtr(set *bool) bool
 
diff --git a/gen/encode.go b/gen/encode.go
index 800c4b19..4e654c26 100644
--- a/gen/encode.go
+++ b/gen/encode.go
@@ -241,6 +241,7 @@ func (e *encodeGen) gMap(m *Map) {
 	e.p.printf("\nfor %s, %s := range %s {", m.Keyidx, m.Validx, vname)
 	e.writeAndCheck(stringTyp, literalFmt, m.Keyidx)
 	e.ctx.PushVar(m.Keyidx)
+	m.Value.SetIsAllowNil(false)
 	next(e, m.Value)
 	e.ctx.Pop()
 	e.p.closeblock()
diff --git a/gen/marshal.go b/gen/marshal.go
index fdc8a48a..59a6e6ec 100644
--- a/gen/marshal.go
+++ b/gen/marshal.go
@@ -245,6 +245,7 @@ func (m *marshalGen) gMap(s *Map) {
 	m.p.printf("\nfor %s, %s := range %s {", s.Keyidx, s.Validx, vname)
 	m.rawAppend(stringTyp, literalFmt, s.Keyidx)
 	m.ctx.PushVar(s.Keyidx)
+	s.Value.SetIsAllowNil(false)
 	next(m, s.Value)
 	m.ctx.Pop()
 	m.p.closeblock()
diff --git a/gen/spec.go b/gen/spec.go
index 93f632d3..d70c9438 100644
--- a/gen/spec.go
+++ b/gen/spec.go
@@ -387,12 +387,7 @@ func (p *printer) rangeBlock(ctx *Context, idx string, iter string, t traversal,
 	ctx.PushVar(idx)
 	// Tags on slices do not extend to the elements, so we always disable allownil on elements.
 	// If we want this to happen in the future, it should be a unique tag.
-	type an interface {
-		SetIsAllowNil(b bool)
-	}
-	if set, ok := inner.(an); ok {
-		set.SetIsAllowNil(false)
-	}
+	inner.SetIsAllowNil(false)
 	p.printf("\n for %s := range %s {", idx, iter)
 	next(t, inner)
 	p.closeblock()
diff --git a/gen/unmarshal.go b/gen/unmarshal.go
index 54de0f2e..f4a9652a 100644
--- a/gen/unmarshal.go
+++ b/gen/unmarshal.go
@@ -269,6 +269,7 @@ func (u *unmarshalGen) gMap(m *Map) {
 	u.p.printf("\nvar %s string; var %s %s; %s--", m.Keyidx, m.Validx, m.Value.TypeName(), sz)
 	u.assignAndCheck(m.Keyidx, stringTyp)
 	u.ctx.PushVar(m.Keyidx)
+	m.Value.SetIsAllowNil(false)
 	next(u, m.Value)
 	u.ctx.Pop()
 	u.p.mapAssign(m)