garble/internal/ssa2ast/polyfill.go

package ssa2ast

import (
	"go/ast"
	"go/token"
	"go/types"
)

func makeMapIteratorPolyfill(tc *typeConverter, mapType *types.Map) (ast.Expr, types.Type, error) {
	keyTypeExpr, err := tc.Convert(mapType.Key())
	if err != nil {
		return nil, nil, err
	}
	valueTypeExpr, err := tc.Convert(mapType.Elem())
	if err != nil {
		return nil, nil, err
	}

	nextType := types.NewSignatureType(nil, nil, nil, nil, types.NewTuple(
		types.NewVar(token.NoPos, nil, "", types.Typ[types.Bool]),
		types.NewVar(token.NoPos, nil, "", mapType.Key()),
		types.NewVar(token.NoPos, nil, "", mapType.Elem()),
	), false)

	// Generated using https://github.com/lu4p/astextract from snippet:
	/*
		func(m map[<key type>]<value type>) func() (bool, <key type>, <value type>) {
			keys := make([]<key type>, 0, len(m))
			for k := range m {
				keys = append(keys, k)
			}
			i := 0
			return func() (ok bool, k <key type>, r <value type>) {
				if i < len(keys) {
					k = keys[i]
					ok, r = true, m[k]
					i++
				}
				return
			}
		}
	*/
	return &ast.FuncLit{
		Type: &ast.FuncType{
			Params: &ast.FieldList{List: []*ast.Field{{
				Names: []*ast.Ident{{Name: "m"}},
				Type: &ast.MapType{
					Key:   keyTypeExpr,
					Value: valueTypeExpr,
				},
			}}},
			Results: &ast.FieldList{List: []*ast.Field{{
				Type: &ast.FuncType{
					Params: &ast.FieldList{},
					Results: &ast.FieldList{List: []*ast.Field{
						{Type: &ast.Ident{Name: "bool"}},
						{Type: keyTypeExpr},
						{Type: valueTypeExpr},
					}},
				},
			}}},
		},
		Body: &ast.BlockStmt{
			List: []ast.Stmt{
				&ast.AssignStmt{
					Lhs: []ast.Expr{&ast.Ident{Name: "keys"}},
					Tok: token.DEFINE,
					Rhs: []ast.Expr{
						&ast.CallExpr{
							Fun: &ast.Ident{Name: "make"},
							Args: []ast.Expr{
								&ast.ArrayType{Elt: keyTypeExpr},
								&ast.BasicLit{Kind: token.INT, Value: "0"},
								&ast.CallExpr{
									Fun:  &ast.Ident{Name: "len"},
									Args: []ast.Expr{&ast.Ident{Name: "m"}},
								},
							},
						},
					},
				},
				&ast.RangeStmt{
					Key: &ast.Ident{Name: "k"},
					Tok: token.DEFINE,
					X:   &ast.Ident{Name: "m"},
					Body: &ast.BlockStmt{
						List: []ast.Stmt{
							&ast.AssignStmt{
								Lhs: []ast.Expr{&ast.Ident{Name: "keys"}},
								Tok: token.ASSIGN,
								Rhs: []ast.Expr{
									&ast.CallExpr{
										Fun: &ast.Ident{Name: "append"},
										Args: []ast.Expr{
											&ast.Ident{Name: "keys"},
											&ast.Ident{Name: "k"},
										},
									},
								},
							},
						},
					},
				},
				&ast.AssignStmt{
					Lhs: []ast.Expr{&ast.Ident{Name: "i"}},
					Tok: token.DEFINE,
					Rhs: []ast.Expr{&ast.BasicLit{Kind: token.INT, Value: "0"}},
				},
				&ast.ReturnStmt{Results: []ast.Expr{
					&ast.FuncLit{
						Type: &ast.FuncType{
							Params: &ast.FieldList{},
							Results: &ast.FieldList{List: []*ast.Field{
								{
									Names: []*ast.Ident{{Name: "ok"}},
									Type:  &ast.Ident{Name: "bool"},
								},
								{
									Names: []*ast.Ident{{Name: "k"}},
									Type:  keyTypeExpr,
								},
								{
									Names: []*ast.Ident{{Name: "r"}},
									Type:  valueTypeExpr,
								},
							}},
						},
						Body: &ast.BlockStmt{
							List: []ast.Stmt{
								&ast.IfStmt{
									Cond: &ast.BinaryExpr{
										X:  &ast.Ident{Name: "i"},
										Op: token.LSS,
										Y: &ast.CallExpr{
											Fun:  &ast.Ident{Name: "len"},
											Args: []ast.Expr{&ast.Ident{Name: "keys"}},
										},
									},
									Body: &ast.BlockStmt{List: []ast.Stmt{
										&ast.AssignStmt{
											Lhs: []ast.Expr{&ast.Ident{Name: "k"}},
											Tok: token.ASSIGN,
											Rhs: []ast.Expr{&ast.IndexExpr{
												X:     &ast.Ident{Name: "keys"},
												Index: &ast.Ident{Name: "i"},
											}},
										},
										&ast.AssignStmt{
											Lhs: []ast.Expr{
												&ast.Ident{Name: "ok"},
												&ast.Ident{Name: "r"},
											},
											Tok: token.ASSIGN,
											Rhs: []ast.Expr{
												&ast.Ident{Name: "true"},
												&ast.IndexExpr{
													X:     &ast.Ident{Name: "m"},
													Index: &ast.Ident{Name: "k"},
												},
											},
										},
										&ast.IncDecStmt{
											X:   &ast.Ident{Name: "i"},
											Tok: token.INC,
										},
									}},
								},
								&ast.ReturnStmt{},
							},
						},
					},
				}},
			},
		},
	}, nextType, nil
}