pruneident_walker.go

package expr

import (
	"fmt"
	"slices"

	"github.com/expr-lang/expr/ast"
	"github.com/expr-lang/expr/file"
	"github.com/expr-lang/expr/parser/operator"
)

type PruneidentWalker struct {
	// idents содержит идентификаторы для удаления в виде map для оптимизации
	idents map[string]struct{}

	// pruned отслеживает местоположения обрезанных узлов для их идентификации
	pruned map[file.Location]struct{}

	// scopes хранит стек областей видимости объявленных переменных
	scopes []scope
}

func NewPruneidentWalker(idents []string) *PruneidentWalker {
	w := &PruneidentWalker{
		idents: make(map[string]struct{}),
	}
	for _, ident := range idents {
		w.idents[ident] = struct{}{}
	}
	return w
}

// scope представляет область видимости для одной переменной.
type scope struct {
	variable string
}

func (w *PruneidentWalker) Walk(node *ast.Node) {
	// очищаем состояние для случая повторного использования
	w.scopes = []scope{}
	w.pruned = make(map[file.Location]struct{})

	w.walk(node)
}

// walk рекурсивно обходит узлы, управляя областями видимости для переменных.
// Реализация аналогична ast.Walk с добавлением обработки областей видимости.
//
//nolint:cyclop,funlen,gocyclo // Вынесение логики в функции ещё больше усложнит восприятие кода.
func (w *PruneidentWalker) walk(node *ast.Node) {
	if node == nil || *node == nil {
		return
	}

	switch n := (*node).(type) {
	case *ast.NilNode:
	case *ast.IdentifierNode:
	case *ast.IntegerNode:
	case *ast.FloatNode:
	case *ast.BoolNode:
	case *ast.StringNode:
	case *ast.ConstantNode:
	case *ast.UnaryNode:
		w.walk(&n.Node)
	case *ast.BinaryNode:
		w.walk(&n.Left)
		w.walk(&n.Right)
	case *ast.ChainNode:
		w.walk(&n.Node)
	case *ast.MemberNode:
		w.walk(&n.Node)
		w.walk(&n.Property)
	case *ast.SliceNode:
		w.walk(&n.Node)
		if n.From != nil {
			w.walk(&n.From)
		}
		if n.To != nil {
			w.walk(&n.To)
		}
	case *ast.CallNode:
		w.walk(&n.Callee)
		for i := range n.Arguments {
			w.walk(&n.Arguments[i])
		}
	case *ast.BuiltinNode:
		for i := range n.Arguments {
			w.walk(&n.Arguments[i])
		}
	case *ast.PredicateNode:
		w.walk(&n.Node)
	case *ast.PointerNode:
	case *ast.VariableDeclaratorNode:
		w.walk(&n.Value)
		w.beginScope(n.Name)
		w.walk(&n.Expr)
		w.endScope()
	case *ast.SequenceNode:
		for i := range n.Nodes {
			w.walk(&n.Nodes[i])
		}
	case *ast.ConditionalNode:
		w.walk(&n.Cond)
		w.walk(&n.Exp1)
		w.walk(&n.Exp2)
	case *ast.ArrayNode:
		for i := range n.Nodes {
			w.walk(&n.Nodes[i])
		}
	case *ast.MapNode:
		for i := range n.Pairs {
			w.walk(&n.Pairs[i])
		}
	case *ast.PairNode:
		w.walk(&n.Key)
		w.walk(&n.Value)
	default:
		panic(fmt.Sprintf("undefined node type (%T)", node))
	}

	w.visit(node)
}

// visit проверяет и модифицирует узлы, удаляя указанные идентификаторы
// и оптимизируя логические выражения после удаления.
//
//nolint:cyclop,funlen,gocyclo,gocognit // Вынесение логики в функции ещё больше усложнит восприятие кода.
func (w *PruneidentWalker) visit(node *ast.Node) {
	switch n := (*node).(type) {
	case *ast.IdentifierNode:
		if w.mustPruned(n.Value) {
			w.pruneNode(node)
		}
	case *ast.UnaryNode:
		if w.prunedNode(n.Node) {
			w.pruneNode(node)
		}
	case *ast.BinaryNode:
		if n.Operator == "in" {
			if in, ok := n.Right.(*ast.IdentifierNode); ok && in.Value == "$env" {
				var sn *ast.StringNode
				sn, ok = n.Left.(*ast.StringNode)
				if ok && w.mustPruned(sn.Value) {
					w.pruneNode(node)
					return
				}
			}
		}

		leftPruned := w.prunedNode(n.Left)
		rightPruned := w.prunedNode(n.Right)

		// Для булевых операторов сохраняем оставшуюся ветку,
		// для остальных - удаляем весь узел при любом удаленном операнде
		if operator.IsBoolean(n.Operator) {
			switch {
			case leftPruned && rightPruned:
				w.pruneNode(node)
			case leftPruned:
				ast.Patch(node, n.Right)
			case rightPruned:
				ast.Patch(node, n.Left)
			}
		} else if leftPruned || rightPruned {
			w.pruneNode(node)
		}
	case *ast.ChainNode:
		if w.prunedNode(n.Node) {
			w.pruneNode(node)
		}
	case *ast.MemberNode:
		if in, ok := n.Node.(*ast.IdentifierNode); ok && in.Value == "$env" {
			var sn *ast.StringNode
			sn, ok = n.Property.(*ast.StringNode)
			if ok && w.mustPruned(sn.Value) {
				w.pruneNode(node)
				return
			}
		}

		if w.prunedNode(n.Node) || w.prunedNode(n.Property) {
			w.pruneNode(node)
		}
	case *ast.SliceNode:
		if w.prunedNode(n.Node) || w.prunedNode(n.From) || w.prunedNode(n.To) {
			w.pruneNode(node)
		}
	case *ast.CallNode:
		if w.prunedNode(n.Callee) {
			w.pruneNode(node)
		}
		for _, arg := range n.Arguments {
			if w.prunedNode(arg) {
				w.pruneNode(node)
			}
		}
	case *ast.BuiltinNode:
		for _, arg := range n.Arguments {
			if w.prunedNode(arg) {
				w.pruneNode(node)
			}
		}
	case *ast.PredicateNode:
		if w.prunedNode(n.Node) {
			w.pruneNode(node)
		}
	case *ast.ConditionalNode:
		if w.prunedNode(n.Cond) || w.prunedNode(n.Exp1) || w.prunedNode(n.Exp2) {
			w.pruneNode(node)
		}
	case *ast.VariableDeclaratorNode:
		if w.prunedNode(n.Value) || w.prunedNode(n.Expr) {
			w.pruneNode(node)
		}
	case *ast.SequenceNode:
		n.Nodes = slices.DeleteFunc(n.Nodes, w.prunedNode)
	case *ast.ArrayNode:
		n.Nodes = slices.DeleteFunc(n.Nodes, w.prunedNode)
	case *ast.MapNode:
		n.Pairs = slices.DeleteFunc(n.Pairs, w.prunedNode)
	case *ast.PairNode:
		if w.prunedNode(n.Key) || w.prunedNode(n.Value) {
			w.pruneNode(node)
		}
	}
}

// mustPruned возвращает true, если идентификатор должен быть удален
// и не находится в текущей области видимости.
func (w *PruneidentWalker) mustPruned(ident string) bool {
	_, exists := w.idents[ident]
	return exists && !w.scoped(ident)
}

// pruneNode заменяет указанный узел на ast.NilNode и сохраняет его позицию.
func (w *PruneidentWalker) pruneNode(node *ast.Node) {
	if node == nil || *node == nil {
		return
	}
	prune := &ast.NilNode{}
	ast.Patch(node, prune)
	(*node).SetType(prune.Type())
	w.pruned[prune.Location()] = struct{}{}
}

// prunedNode проверяет, был ли узел помечен как удаленный.
func (w *PruneidentWalker) prunedNode(node ast.Node) bool {
	if node == nil {
		return false
	}
	if n, ok := node.(*ast.NilNode); ok {
		_, exists := w.pruned[n.Location()]
		return exists
	}
	return false
}

// beginScope добавляет новую область видимости для переменной.
func (w *PruneidentWalker) beginScope(variable string) {
	w.scopes = append(w.scopes, scope{variable: variable})
}

// endScope удаляет последнюю добавленную область видимости.
func (w *PruneidentWalker) endScope() {
	w.scopes = w.scopes[:len(w.scopes)-1]
}

// scoped проверяет, существует ли переменная в текущей области видимости.
func (w *PruneidentWalker) scoped(variable string) bool {
	for i := len(w.scopes) - 1; i >= 0; i-- {
		if w.scopes[i].variable == variable {
			return true
		}
	}
	return false
}