// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The deadlocals pass removes assignments to unused local variables.
package deadlocals

import (
	"cmd/compile/internal/base"
	"cmd/compile/internal/ir"
	"cmd/compile/internal/types"
	"cmd/internal/src"
	"fmt"
	"go/constant"
)

// Funcs applies the deadlocals pass to fns.
func Funcs(fns []*ir.Func) {
	if base.Flag.N != 0 || base.Debug.NoDeadLocals != 0 {
		return
	}

	zero := ir.NewBasicLit(base.AutogeneratedPos, types.Types[types.TINT], constant.MakeInt64(0))

	for _, fn := range fns {
		if fn.IsClosure() {
			continue
		}

		v := newVisitor(fn)
		v.nodes(fn.Body)

		for _, k := range v.defsKeys {
			assigns := v.defs[k]
			for _, as := range assigns {
				// Kludge for "missing func info" linker panic.
				// See also closureInitLSym in inline/inl.go.
				if clo, ok := (*as.rhs).(*ir.ClosureExpr); ok && clo.Op() == ir.OCLOSURE {
					if clo.Func.IsClosure() {
						ir.InitLSym(clo.Func, true)
					}
				}

				*as.lhs = ir.BlankNode
				*as.rhs = zero
			}
		}
	}
}

type visitor struct {
	curfn *ir.Func
	// defs[name] contains assignments that can be discarded if name can be discarded.
	// if defs[name] is defined nil, then name is actually used.
	defs     map[*ir.Name][]assign
	defsKeys []*ir.Name // insertion order of keys, for reproducible iteration (and builds)

	doNode func(ir.Node) bool
}

type assign struct {
	pos      src.XPos
	lhs, rhs *ir.Node
}

func newVisitor(fn *ir.Func) *visitor {
	v := &visitor{
		curfn: fn,
		defs:  make(map[*ir.Name][]assign),
	}
	v.doNode = func(n ir.Node) bool {
		v.node(n)
		return false
	}
	return v
}

func (v *visitor) node(n ir.Node) {
	if n == nil {
		return
	}

	switch n.Op() {
	default:
		ir.DoChildrenWithHidden(n, v.doNode)
	case ir.OCLOSURE:
		n := n.(*ir.ClosureExpr)
		v.nodes(n.Init())
		for _, cv := range n.Func.ClosureVars {
			v.node(cv)
		}
		v.nodes(n.Func.Body)

	case ir.ODCL:
		// ignore
	case ir.ONAME:
		n := n.(*ir.Name)
		n = n.Canonical()
		if isLocal(n, false) {
			// Force any lazy definitions.
			s, ok := v.defs[n]
			if !ok {
				v.defsKeys = append(v.defsKeys, n)
			}
			v.defs[n] = nil
			for _, as := range s {
				// do the visit that was skipped in v.assign when as was appended to v.defs[n]
				v.node(*as.rhs)
			}
		}

	case ir.OAS:
		n := n.(*ir.AssignStmt)
		v.assign(n.Pos(), &n.X, &n.Y, false)
	case ir.OAS2:
		n := n.(*ir.AssignListStmt)

		// If all LHS vars are blank, treat them as intentional
		// uses of corresponding RHS vars.  If any are non-blank
		// then any blanks are discards.
		hasNonBlank := false
		for i := range n.Lhs {
			if !ir.IsBlank(n.Lhs[i]) {
				hasNonBlank = true
				break
			}
		}
		for i := range n.Lhs {
			v.assign(n.Pos(), &n.Lhs[i], &n.Rhs[i], hasNonBlank)
		}
	}
}

func (v *visitor) nodes(list ir.Nodes) {
	for _, n := range list {
		v.node(n)
	}
}

func hasEffects(n ir.Node) bool {
	if n == nil {
		return false
	}
	if len(n.Init()) != 0 {
		return true
	}

	switch n.Op() {
	// TODO(mdempsky): More.
	case ir.ONAME, ir.OLITERAL, ir.ONIL, ir.OCLOSURE:
		return false
	}
	return true
}

func (v *visitor) assign(pos src.XPos, lhs, rhs *ir.Node, blankIsNotUse bool) {
	name, ok := (*lhs).(*ir.Name)
	if !ok {
		v.node(*lhs) // XXX: Interpret as variable, not value.
		v.node(*rhs)
		return
	}
	name = name.Canonical()

	if isLocal(name, blankIsNotUse) && !hasEffects(*rhs) {
		if s, ok := v.defs[name]; !ok || s != nil {
			// !ok || s != nil is FALSE if previously "v.defs[name] = nil" -- that marks a use.
			if !ok {
				v.defsKeys = append(v.defsKeys, name)
			}
			v.defs[name] = append(s, assign{pos, lhs, rhs})
			return // don't visit rhs unless that node ends up live, later.
		}
	}

	v.node(*rhs)
}

func isLocal(n *ir.Name, blankIsNotUse bool) bool {
	if ir.IsBlank(n) {
		// Treat single assignments as intentional use (false), anything else is a discard (true).
		return blankIsNotUse
	}

	switch n.Class {
	case ir.PAUTO, ir.PPARAM:
		return true
	case ir.PPARAMOUT:
		return false
	case ir.PEXTERN, ir.PFUNC:
		return false
	}
	panic(fmt.Sprintf("unexpected Class: %+v", n))
}