// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package analyze

import (
	"errors"
	"fmt"
	"iter"
	"slices"
	"strings"

	"github.com/cilium/ebpf/asm"
)

// leaderKey is used to store the leader metadata in an instruction's metadata.
type leaderKey struct{}

// A leader is an instruction at the beginning of a basic block.
type leader struct {
	// predecessors are instructions in other blocks that are always executed
	// before this instruction.
	predecessors []*asm.Instruction

	// block is the block that this instruction is the start of.
	block *Block
}

// setLeaderMeta sets the leader metadata for an instruction. This metadata
// is used to mark the start of a basic block and to store information about
// the block and its predecessors.
func setLeaderMeta(ins *asm.Instruction, meta *leader) {
	ins.Metadata.Set(leaderKey{}, meta)
}

// getLeaderMeta retrieves the leader metadata for an instruction.
func getLeaderMeta(ins *asm.Instruction) *leader {
	val := ins.Metadata.Get(leaderKey{})
	meta, ok := val.(*leader)
	if !ok {
		return nil
	}
	return meta
}

// addPredecessors adds one or more predecessor instructions to the list of
// predecessors for the given instruction. If a predecessor is already in the
// list, it is not added again.
//
// This is used to track the control flow graph of the program, where each
// instruction can have multiple predecessors (i.e. it can be reached from
// multiple branches). Initializes the instruction's leader metadata if it does
// not exist yet.
func addPredecessors(ins *asm.Instruction, preds ...*asm.Instruction) {
	l := getLeaderMeta(ins)
	if l == nil {
		l = &leader{}
		setLeaderMeta(ins, l)
	}
	for _, pred := range preds {
		if pred == nil {
			continue
		}
		if !slices.Contains(l.predecessors, pred) {
			l.predecessors = append(l.predecessors, pred)
		}
	}
}

// edgeKey is used to store the edge metadata in an instruction's metadata.
type edgeKey struct{}

// edge is a metadata structure that is associated with an instruction marking
// the end of a basic block. It can have a branch target (the target of a jump
// instruction) and a fallthrough target (the next instruction in the
// instruction stream that is executed if the branch is not taken).
type edge struct {
	branch   *asm.Instruction
	fthrough *asm.Instruction

	block *Block
}

// setEdgeMeta sets the edge metadata for an instruction.
func setEdgeMeta(ins *asm.Instruction, meta *edge) {
	ins.Metadata.Set(edgeKey{}, meta)
}

// getEdgeMeta retrieves the edge metadata for an instruction.
func getEdgeMeta(ins *asm.Instruction) *edge {
	val := ins.Metadata.Get(edgeKey{})
	meta, ok := val.(*edge)
	if !ok {
		return nil
	}
	return meta
}

// setEdgeBranchTarget sets the branch target for an edge instruction. This is
// used to mark the target of a jump instruction that branches to another basic
// block.
func setEdgeBranchTarget(ins *asm.Instruction, target *asm.Instruction) {
	e := getEdgeMeta(ins)
	if e == nil {
		e = &edge{}
		setEdgeMeta(ins, e)
	}
	e.branch = target
}

// setEdgeFallthrough sets the fallthrough target for an edge instruction. This
// is used to mark the next instruction in the instruction stream that is
// executed if the branch is not taken, typically the instruction immediately
// following the branch instruction.
func setEdgeFallthrough(ins *asm.Instruction, target *asm.Instruction) {
	if ins == nil {
		return
	}
	e := getEdgeMeta(ins)
	if e == nil {
		e = &edge{}
		setEdgeMeta(ins, e)
	}
	e.fthrough = target
}

// setBranchTarget creates a two-way association between both the branch
// instruction and its target instruction, as well as the target instruction and
// its natural predecessor. prev may be nil if the branch target is the first
// instruction in the program.
//
// This process creates two edges and a leader, updating existing metadata if
// the instructions were already marked as leaders or edges.
func setBranchTarget(branch, target, prev *asm.Instruction) {
	// Associate the branch instruction with its target. The target becomes a
	// leader, the branch instruction becomes an edge.
	setEdgeBranchTarget(branch, target)

	// Create a reverse link from the branch target to both the branch (jump)
	// instructions and the target's predecessor.
	if canFallthrough(prev) {
		// Creating a leader implicitly means making the instruction before it an
		// edge with a fallthrough target. Associate the target with its predecessor.
		setEdgeFallthrough(prev, target)
		addPredecessors(target, branch, prev)
	} else {
		// If the instruction preceding the branch target cannot fall through (Ja,
		// Exit), don't register it as a predecessor.
		setEdgeFallthrough(prev, nil)
		addPredecessors(target, branch)
	}
}

// A Block is a contiguous sequence of instructions that are executed together.
// Boundaries are defined by branching instructions.
//
// Blocks are attached to instructions via metadata and should not be modified
// after being created.
//
// It should never contain direct references to the original asm.Instructions
// since copying the ProgramSpec won't update pointers to the new copied insns.
// This is a problem when modifying instructions through
// [Blocks.LiveInstructions] after reachability analysis, since it would modify
// the original ProgramSpec's instructions.
type Block struct {
	id         uint64
	raw        asm.RawInstructionOffset
	start, end int

	predecessors []*Block
	branch       *Block
	fthrough     *Block
}

func (b *Block) leader(insns asm.Instructions) *leader {
	if len(insns) == 0 {
		return nil
	}
	return getLeaderMeta(&insns[b.start])
}

func (b *Block) edge(insns asm.Instructions) *edge {
	if len(insns) == 0 {
		return nil
	}

	if b.end >= len(insns) {
		return nil
	}

	return getEdgeMeta(&insns[b.end])
}

func (b *Block) iterateLocal(insns asm.Instructions) *BlockIterator {
	if b.start < 0 || b.end < 0 || b.start >= len(insns) || b.end >= len(insns) {
		return nil
	}

	return &BlockIterator{
		block:  b,
		insns:  insns,
		index:  b.start,
		offset: b.raw,
		local:  true,
	}
}

func (b *Block) iterateGlobal(blocks Blocks, insns asm.Instructions) *BlockIterator {
	if b.start < 0 || b.end < 0 || b.start >= len(insns) || b.end >= len(insns) {
		return nil
	}
	return &BlockIterator{
		blocks: blocks,
		block:  b,
		insns:  insns,
		index:  b.start,
		offset: b.raw,
		local:  false,
	}
}

// predecessor returns the previous Block in the control flow if there is
// exactly one predecessor, otherwise returns nil.
//
// This is useful for walking the control flow backwards when there is no
// branching, e.g. for finding the last instruction that wrote to a register
// before it is read.
func (b *Block) predecessor() *Block {
	if len(b.predecessors) == 1 {
		return b.predecessors[0]
	}
	return nil
}

func (b *Block) String() string {
	return b.Dump(nil)
}

func (b *Block) Dump(insns asm.Instructions) string {
	var sb strings.Builder
	sb.WriteString("Predecessors: [")
	for i, from := range b.predecessors {
		if i > 0 {
			sb.WriteString(", ")
		}
		sb.WriteString(fmt.Sprintf("%d", from.id))
	}
	sb.WriteString("]\n")
	sb.WriteString(fmt.Sprintf("Start: %d (raw %d), end: %d\n", b.start, b.raw, b.end))
	sb.WriteString("\n")

	if len(insns) != 0 {
		sb.WriteString("Instructions:\n")
		i := b.iterateLocal(insns)
		for i.Next() {
			ins := i.Instruction()
			if ins.Symbol() != "" {
				fmt.Fprintf(&sb, "\t%s:\n", ins.Symbol())
			}
			if src := ins.Source(); src != nil {
				line := strings.TrimSpace(src.String())
				if line != "" {
					fmt.Fprintf(&sb, "\t%*s; %s\n", 4, " ", line)
				}
			}
			fmt.Fprintf(&sb, "\t%*d: %v\n", 4, i.Offset(), ins)
		}
		sb.WriteString("\n")
	} else {
		sb.WriteString("Instructions: not provided, call Dump() with insns\n")
	}

	if b.branch != nil {
		sb.WriteString("Branch: ")
		sb.WriteString(fmt.Sprintf("%d", b.branch.id))
		sb.WriteString("\n")
	}

	if b.fthrough != nil {
		sb.WriteString("Fallthrough: ")
		sb.WriteString(fmt.Sprintf("%d", b.fthrough.id))
		sb.WriteString("\n")
	}

	return sb.String()
}

// maxDepth is the maximum depth of block traversal when backtracking to
// predecessors. Used in favor of a visited set since it's much cheaper than
// frequent map lookups. Typical depth while looking for map pointer loads is
// 1-3 with some double-digit outliers.
const maxDepth = 128

// BlockIterator is an iterator over the instructions in a block or a list of
// blocks.
//
// It can be configured to iterate locally within a block or globally across
// multiple blocks. When iterating globally, it will roll over to the next or
// previous block when reaching the end or start of the current block,
// respectively.
//
// The iterator tracks the raw instruction offset of the current instruction
// when iterating forwards, but not when iterating backwards since that would
// require summing up instruction sizes from the start of the block. Raw offsets
// are only used for dumping instructions in forward order.
type BlockIterator struct {
	blocks Blocks
	block  *Block

	insns asm.Instructions
	depth uint8

	ins   *asm.Instruction
	index int

	// offset is zero when backtracking to predecessors or when iterating
	// backwards from the end of a block (e.g. with a new iterator).
	offset asm.RawInstructionOffset

	local bool
}

func (i *BlockIterator) Instruction() *asm.Instruction {
	return i.ins
}

func (i *BlockIterator) Index() int {
	return i.index
}

func (i *BlockIterator) Offset() asm.RawInstructionOffset {
	return i.offset
}

func (i *BlockIterator) Clone() *BlockIterator {
	return &BlockIterator{
		i.blocks,
		i.block,
		i.insns,
		i.depth,
		i.ins,
		i.index,
		i.offset,
		i.local,
	}
}

// nextBlock pulls the next block by identifier, if it exists. Otherwise,
// returns false.
//
// Positions the iterator at the start of the next block. Offset is updated to
// the raw offset of the first instruction in the next block.
func (i *BlockIterator) nextBlock() bool {
	if i.block == nil {
		return false
	}

	if i.block.id+1 >= i.blocks.count() {
		return false
	}
	next := i.blocks[i.block.id+1]

	// Reset loop detection when moving forward since we're primarily concerned
	// with infinite loops while backtracking continuously. Fallthroughs always
	// point to the end of the program eventually.
	i.depth = 0

	i.block = next
	i.index = i.block.start
	i.offset = i.block.raw
	i.ins = &i.insns[i.index]

	return true
}

// prevBlock pulls the previous block in the control flow if there is exactly
// one predecessor, otherwise returns false.
//
// Sometimes, map pointers are loaded into a register in a previous block.
// Backtracking into multiple predecessors is not useful since the contents of
// the pointer register would be ambiguous if assigned from multiple
// predecessors.
//
// Positions the iterator at the end of the previous block. Offset is set to 0
// since backtracking requires summing up instruction sizes from the start of
// the block.
func (i *BlockIterator) prevBlock() bool {
	if i.block == nil {
		return false
	}

	prev := i.block.predecessor()
	if prev == nil || prev == i.block {
		return false
	}

	// Simple loop detection to avoid infinite loops when backtracking
	// continuously. Depth gets reset when moving forward.
	if i.depth >= maxDepth {
		return false
	}
	i.depth++

	i.block = prev
	i.index = i.block.end

	// Raw offset tracking disabled for backtracking since it requires summing
	// up instruction sizes from the start of the block. Raw offsets are only
	// used for dumping instructions in forward order.
	i.offset = 0

	i.ins = &i.insns[i.index]

	return true
}

// Next advances the iterator to the next instruction in the block. If the end
// of the block is reached, it will either stop (if iterating locally) or roll
// over to the next block (if iterating globally).
func (i *BlockIterator) Next() bool {
	if i.block == nil || i.index < i.block.start || i.index > i.block.end {
		return false
	}

	// First call to Next with this iterator, pull the first insn and return.
	if i.ins == nil {
		i.ins = &i.insns[i.index]
		return true
	}

	if i.index+1 > i.block.end {
		if !i.local {
			// Iterating globally, roll over to the next block if it exists.
			return i.nextBlock()
		}

		// Iterating locally, stop here.
		return false
	}

	i.index++
	i.offset += asm.RawInstructionOffset(i.ins.Size() / asm.InstructionSize)
	i.ins = &i.insns[i.index]

	return true
}

func (i *BlockIterator) Previous() bool {
	if i.block == nil || i.index < i.block.start || i.index > i.block.end {
		return false
	}

	if i.ins == nil {
		i.index = i.block.end
		i.ins = &i.insns[i.index]

		// Raw offset tracking disabled for backtracking since it requires summing
		// up instruction sizes from the start of the block. Raw offsets are only
		// used for dumping instructions in forward order.
		i.offset = 0

		return true
	}

	if i.index-1 < i.block.start {
		if !i.local {
			// Iterating globally, roll over to the previous block if it exists.
			return i.prevBlock()
		}

		// Iterating locally, stop here.
		return false
	}

	i.index--
	i.ins = &i.insns[i.index]

	// Prevent offset underflow.
	raw := asm.RawInstructionOffset(i.ins.Size() / asm.InstructionSize)
	i.offset = min(i.offset, i.offset-raw)

	return true
}

// getBlock retrieves the block associated with an instruction. It checks both
// the leader and edge metadata to find the block. If neither is found, it
// returns nil, indicating that the instruction forms neither the start nor end
// of a basic block.
func getBlock(ins *asm.Instruction) *Block {
	l := getLeaderMeta(ins)
	if l != nil {
		return l.block
	}

	e := getEdgeMeta(ins)
	if e != nil {
		return e.block
	}

	return nil
}

// Blocks is a list of basic blocks.
type Blocks []*Block

func (bl Blocks) count() uint64 {
	return uint64(len(bl))
}

func (bl *Blocks) add(b *Block) {
	if b == nil {
		return
	}

	b.id = uint64(bl.count())
	*bl = append(*bl, b)
}

func (bl Blocks) first() *Block {
	if len(bl) == 0 {
		return nil
	}
	return bl[0]
}

func (bl Blocks) last() *Block {
	if len(bl) == 0 {
		return nil
	}
	return bl[len(bl)-1]
}

func (bl Blocks) iterate(insns asm.Instructions) *BlockIterator {
	if len(bl) == 0 {
		return nil
	}
	return bl.first().iterateGlobal(bl, insns)
}

func (bl Blocks) String() string {
	return bl.Dump(nil)
}

func (bl Blocks) Dump(insns asm.Instructions) string {
	var sb strings.Builder
	for _, block := range bl {
		sb.WriteString(fmt.Sprintf("\n=== Block %d ===\n", block.id))
		sb.WriteString(block.Dump(insns))
		sb.WriteString("\n")
	}
	return sb.String()
}

// MakeBlocks returns a list of basic blocks of instructions that are always
// executed together. Multiple calls on the same insns will return the same
// Blocks object.
//
// Blocks are created by finding branches and jump targets in the given insns
// and cutting up the instruction stream accordingly.
func MakeBlocks(insns asm.Instructions) (Blocks, error) {
	if len(insns) == 0 {
		return nil, errors.New("insns is empty, cannot compute blocks")
	}

	if blocks := loadBlocks(insns); blocks != nil {
		return blocks, nil
	}

	blocks, err := computeBlocks(insns)
	if err != nil {
		return nil, fmt.Errorf("computing blocks: %w", err)
	}

	if err := storeBlocks(insns, blocks); err != nil {
		return nil, fmt.Errorf("storing blocks: %w", err)
	}

	return blocks, nil
}

// computeBlocks computes the basic blocks from the given instruction stream.
func computeBlocks(insns asm.Instructions) (Blocks, error) {
	targets, err := rawJumpTargets(insns)
	if err != nil {
		return nil, fmt.Errorf("collecting jump targets: %w", err)
	}

	if err := tagLeadersAndEdges(insns, targets); err != nil {
		return nil, fmt.Errorf("tagging instructions: %w", err)
	}

	blocks, err := allocateBlocks(insns)
	if err != nil {
		return nil, fmt.Errorf("allocating blocks: %w", err)
	}

	if err := connectBlocks(blocks, insns); err != nil {
		return nil, fmt.Errorf("connecting blocks: %w", err)
	}

	return blocks, nil
}

// blocksKey is used to store Blocks in an instruction's metadata.
type blocksKey struct{}

// storeBlocks associates the given Blocks with the first instruction in the
// given insns.
//
// If insns is empty, does nothing.
func storeBlocks(insns asm.Instructions, bl Blocks) error {
	if len(insns) == 0 {
		return errors.New("insns is empty, cannot store Blocks")
	}

	insns[0].Metadata.Set(blocksKey{}, bl)

	return nil
}

// loadBlocks retrieves the Blocks associated with the first instruction in the
// given insns.
//
// If no Blocks is present, returns nil.
func loadBlocks(insns asm.Instructions) Blocks {
	if len(insns) == 0 {
		return nil
	}

	blocks, ok := insns[0].Metadata.Get(blocksKey{}).(Blocks)
	if !ok {
		return nil
	}

	return blocks
}

// rawJumpTargets returns a map of raw instruction offsets to jump targets,
// where each target is a logical instruction in the instruction stream.
//
// The raw instruction offsets are the offsets of the instructions in the raw
// bytecode, which may not correspond to the logical instruction indices due to
// variable instruction sizes (e.g. dword loads).
func rawJumpTargets(insns asm.Instructions) (rawTargets, error) {
	// Jump offsets are in raw instructions of size [asm.InstructionSize], but
	// some instructions are 2x the size of a normal instruction (e.g. dword
	// loads). Find the raw offsets of all jump targets and mark them for
	// resolution.
	targets := make(rawTargets)
	i := insns.Iterate()
	for i.Next() {
		target, ok := jumpTarget(i.Offset, i.Ins)
		if !ok {
			continue
		}

		// Queue the target as a 'raw' leader to be resolved to a logical insn in
		// the next loop.
		targets.add(target)

		// Mark the instruction as an incomplete edge to avoid re-checking if each
		// insn is a jump in a subsequent step.
		setEdgeMeta(i.Ins, &edge{})
	}

	if len(targets) == 0 {
		// No jump targets to resolve.
		return nil, nil
	}

	// Second loop for finding the [asm.Instruction] for each raw offset
	// identified in the previous step.
	next, stop := iter.Pull(targets.keysSorted())
	defer stop()

	// Memoize the next leader so we don't need a map lookup for every insn.
	nextTarget, ok := next()
	if !ok {
		return nil, errors.New("no jump target to resolve, this is a bug")
	}

	i = insns.Iterate()
	for i.Next() {
		if i.Offset != nextTarget {
			continue
		}

		// Map the raw instruction offset to its logical instruction.
		targets.resolve(i.Offset, i.Index, i.Ins)

		// Pull the next target to resolve.
		nextTarget, ok = next()
		if !ok {
			// No more targets to resolve.
			break
		}
	}

	return targets, nil
}

// tagLeadersAndEdges tags the instructions in the given instruction stream
// as leaders and/or edges based on their control flow properties. It identifies
// the first instruction as a leader without predecessors, the last instruction
// as an edge without a branch or fallthrough, and processes jump instructions
// to create leaders for their targets and edges for their predecessors.
//
// Returns error if any edge instruction does not have a target instruction at
// the specified raw offset.
func tagLeadersAndEdges(insns asm.Instructions, targets rawTargets) error {
	// Mark first insn as leader without predecessors, last insn as an edge
	// without a branch or fallthrough.
	setLeaderMeta(&insns[0], &leader{})
	setEdgeMeta(&insns[len(insns)-1], &edge{})

	if len(targets) == 0 {
		// No jump targets to resolve.
		return nil
	}

	// Find all jump instructions, create leaders for their targets and edges for
	// their predecessors.
	i := insns.Iterate()
	for i.Next() {
		// If the insn was identified as an edge in a prior step, add it as a
		// predecessor to the next instruction and to the branch target.
		e := getEdgeMeta(i.Ins)
		if e == nil {
			continue
		}

		// If the instruction is a branch, we need to find the target instruction
		// and set it as the branch target.
		raw, ok := jumpTarget(i.Offset, i.Ins)
		if !ok {
			// Edge doesn't have a jump target. This could be an exit or call
			// instruction, in which case there's no jump target to resolve and no
			// leader to create.
			continue
		}

		tgt := targets.get(raw)
		if tgt == nil {
			return fmt.Errorf("edge %v has no target instruction at offset %d", i.Ins, raw)
		}

		// In case of a jump to the first instruction, the target has no
		// predecessor, so we need a bounds check.
		var prev *asm.Instruction
		if tgt.index-1 >= 0 {
			prev = &insns[tgt.index-1]
		}
		setBranchTarget(i.Ins, tgt.ins, prev)

		// No next instruction, don't set a fallthrough target.
		if i.Index == len(insns)-1 {
			continue
		}

		// If the instruction is an unconditional jump, don't consider the next
		// instruction a fallthrough target.
		if i.Ins.OpCode.JumpOp() == asm.Ja {
			continue
		}

		next := &insns[i.Index+1]
		addPredecessors(next, i.Ins)
		setEdgeFallthrough(i.Ins, next)
	}

	return nil
}

// allocateBlocks returns a list of blocks based on leaders and edges identified
// in prior stages. It creates a new block whenever it encounters a leader
// instruction and finalizes the current one when it reaches an edge
// instruction. No blocks are pointing to each other yet, this is done in a
// subsequent step.
func allocateBlocks(insns asm.Instructions) (Blocks, error) {
	// Expect at least one block.
	blocks := make(Blocks, 0, 1)

	var block *Block
	i := insns.Iterate()
	for i.Next() {
		// Roll over to the next block if this is a leader.
		if nextBlock := maybeAllocateBlock(i); nextBlock != nil {
			block = nextBlock
			blocks.add(block)
		}

		// Finalize the block if we've reached an edge.
		maybeFinalizeBlock(block, i)
	}

	if blocks.count() == 0 {
		return nil, errors.New("no blocks created, this is a bug")
	}

	return blocks, nil
}

// maybeAllocateBlock allocates a new block for the instruction pointed to by
// the iterator if it is a leader instruction. If the instruction is not a
// leader, it returns nil. This is used to start a new basic block when
// encountering a leader instruction in the instruction stream.
func maybeAllocateBlock(i *asm.InstructionIterator) *Block {
	l := getLeaderMeta(i.Ins)
	if l == nil {
		return nil
	}
	l.block = &Block{
		start: i.Index,
		raw:   i.Offset,
	}
	return l.block
}

// maybeFinalizeBlock finalizes the current block by populating its insns field
// and associating it with the edge metadata if the instruction is an edge.
// If the instruction is not an edge or the given block is nil, it does nothing.
func maybeFinalizeBlock(blk *Block, i *asm.InstructionIterator) {
	e := getEdgeMeta(i.Ins)
	if e == nil {
		return
	}
	if blk == nil {
		return
	}
	blk.end = i.Index
	e.block = blk
}

// connectBlocks connects the blocks in the given block list by setting their
// predecessors, branch and fallthrough targets based on the relationships
// between instructions identified in prior steps. Assumes that blocks have been
// allocated and that the leaders and edges have been tagged.
func connectBlocks(blocks Blocks, insns asm.Instructions) error {
	if blocks.count() == 0 {
		return errors.New("no blocks to connect, this is a bug")
	}

	// Wire all blocks together by setting their predecessors, branch and
	// fallthrough targets.
	for _, blk := range blocks {
		// Predecessors of the first instruction are the block's predecessors.
		leader := blk.leader(insns)
		if leader == nil {
			return fmt.Errorf("block %d has no leader", blk.id)
		}

		blk.predecessors = make([]*Block, 0, len(leader.predecessors))
		for _, pi := range leader.predecessors {
			b := getBlock(pi)
			if b == nil {
				return fmt.Errorf("predecessor instruction %v has no block", pi)
			}
			blk.predecessors = append(blk.predecessors, b)
		}

		// Branch/fthrough targets of the last instruction are the block's branch
		// and fallthrough targets.
		edge := blk.edge(insns)
		if edge == nil {
			return fmt.Errorf("block %d has no edge", blk.id)
		}

		if edge.branch != nil {
			// If the edge has a branch target, set it as the block's branch target.
			b := getBlock(edge.branch)
			if b == nil {
				return fmt.Errorf("branch target %v has no block", edge.branch)
			}
			blk.branch = b
		}

		if edge.fthrough != nil {
			// If the edge has a fallthrough target, set it as the block's fallthrough
			// target.
			b := getBlock(edge.fthrough)
			if b == nil {
				return fmt.Errorf("fallthrough target %v has no block", edge.fthrough)
			}
			blk.fthrough = b
		}
	}

	return nil
}
