// Copyright 2022 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.

// A note on line numbers: when working with line numbers, we always use the
// binary-visible relative line number. i.e., the line number as adjusted by
// //line directives (ctxt.InnermostPos(ir.Node.Pos()).RelLine()). Use
// NodeLineOffset to compute line offsets.
//
// If you are thinking, "wait, doesn't that just make things more complex than
// using the real line number?", then you are 100% correct. Unfortunately,
// pprof profiles generated by the runtime always contain line numbers as
// adjusted by //line directives (because that is what we put in pclntab). Thus
// for the best behavior when attempting to match the source with the profile
// it makes sense to use the same line number space.
//
// Some of the effects of this to keep in mind:
//
//  - For files without //line directives there is no impact, as RelLine() ==
//    Line().
//  - For functions entirely covered by the same //line directive (i.e., a
//    directive before the function definition and no directives within the
//    function), there should also be no impact, as line offsets within the
//    function should be the same as the real line offsets.
//  - Functions containing //line directives may be impacted. As fake line
//    numbers need not be monotonic, we may compute negative line offsets. We
//    should accept these and attempt to use them for best-effort matching, as
//    these offsets should still match if the source is unchanged, and may
//    continue to match with changed source depending on the impact of the
//    changes on fake line numbers.
//  - Functions containing //line directives may also contain duplicate lines,
//    making it ambiguous which call the profile is referencing. This is a
//    similar problem to multiple calls on a single real line, as we don't
//    currently track column numbers.
//
// Long term it would be best to extend pprof profiles to include real line
// numbers. Until then, we have to live with these complexities. Luckily,
// //line directives that change line numbers in strange ways should be rare,
// and failing PGO matching on these files is not too big of a loss.

package pgo

import (
	"cmd/compile/internal/base"
	"cmd/compile/internal/ir"
	"cmd/compile/internal/pgo/internal/graph"
	"cmd/compile/internal/typecheck"
	"cmd/compile/internal/types"
	"fmt"
	"internal/profile"
	"os"
)

// IRGraph is a call graph with nodes pointing to IRs of functions and edges
// carrying weights and callsite information.
//
// Nodes for indirect calls may have missing IR (IRNode.AST == nil) if the node
// is not visible from this package (e.g., not in the transitive deps). Keeping
// these nodes allows determining the hottest edge from a call even if that
// callee is not available.
//
// TODO(prattmic): Consider merging this data structure with Graph. This is
// effectively a copy of Graph aggregated to line number and pointing to IR.
type IRGraph struct {
	// Nodes of the graph
	IRNodes map[string]*IRNode
}

// IRNode represents a node (function) in the IRGraph.
type IRNode struct {
	// Pointer to the IR of the Function represented by this node.
	AST *ir.Func
	// Linker symbol name of the Function represented by this node.
	// Populated only if AST == nil.
	LinkerSymbolName string

	// Set of out-edges in the callgraph. The map uniquely identifies each
	// edge based on the callsite and callee, for fast lookup.
	OutEdges map[NodeMapKey]*IREdge
}

// Name returns the symbol name of this function.
func (i *IRNode) Name() string {
	if i.AST != nil {
		return ir.LinkFuncName(i.AST)
	}
	return i.LinkerSymbolName
}

// IREdge represents a call edge in the IRGraph with source, destination,
// weight, callsite, and line number information.
type IREdge struct {
	// Source and destination of the edge in IRNode.
	Src, Dst       *IRNode
	Weight         int64
	CallSiteOffset int // Line offset from function start line.
}

// NodeMapKey represents a hash key to identify unique call-edges in profile
// and in IR. Used for deduplication of call edges found in profile.
//
// TODO(prattmic): rename to something more descriptive.
type NodeMapKey struct {
	CallerName     string
	CalleeName     string
	CallSiteOffset int // Line offset from function start line.
}

// Weights capture both node weight and edge weight.
type Weights struct {
	NFlat   int64
	NCum    int64
	EWeight int64
}

// CallSiteInfo captures call-site information and its caller/callee.
type CallSiteInfo struct {
	LineOffset int // Line offset from function start line.
	Caller     *ir.Func
	Callee     *ir.Func
}

// Profile contains the processed PGO profile and weighted call graph used for
// PGO optimizations.
type Profile struct {
	// Aggregated NodeWeights and EdgeWeights across the profile. This
	// helps us determine the percentage threshold for hot/cold
	// partitioning.
	TotalNodeWeight int64
	TotalEdgeWeight int64

	// NodeMap contains all unique call-edges in the profile and their
	// aggregated weight.
	NodeMap map[NodeMapKey]*Weights

	// WeightedCG represents the IRGraph built from profile, which we will
	// update as part of inlining.
	WeightedCG *IRGraph
}

// New generates a profile-graph from the profile.
func New(profileFile string) (*Profile, error) {
	f, err := os.Open(profileFile)
	if err != nil {
		return nil, fmt.Errorf("error opening profile: %w", err)
	}
	defer f.Close()
	profile, err := profile.Parse(f)
	if err != nil {
		return nil, fmt.Errorf("error parsing profile: %w", err)
	}

	if len(profile.Sample) == 0 {
		// We accept empty profiles, but there is nothing to do.
		return nil, nil
	}

	valueIndex := -1
	for i, s := range profile.SampleType {
		// Samples count is the raw data collected, and CPU nanoseconds is just
		// a scaled version of it, so either one we can find is fine.
		if (s.Type == "samples" && s.Unit == "count") ||
			(s.Type == "cpu" && s.Unit == "nanoseconds") {
			valueIndex = i
			break
		}
	}

	if valueIndex == -1 {
		return nil, fmt.Errorf(`profile does not contain a sample index with value/type "samples/count" or cpu/nanoseconds"`)
	}

	g := graph.NewGraph(profile, &graph.Options{
		SampleValue: func(v []int64) int64 { return v[valueIndex] },
	})

	p := &Profile{
		NodeMap: make(map[NodeMapKey]*Weights),
		WeightedCG: &IRGraph{
			IRNodes: make(map[string]*IRNode),
		},
	}

	// Build the node map and totals from the profile graph.
	if err := p.processprofileGraph(g); err != nil {
		return nil, err
	}

	if p.TotalNodeWeight == 0 || p.TotalEdgeWeight == 0 {
		return nil, nil // accept but ignore profile with no samples.
	}

	// Create package-level call graph with weights from profile and IR.
	p.initializeIRGraph()

	return p, nil
}

// processprofileGraph builds various maps from the profile-graph.
//
// It initializes NodeMap and Total{Node,Edge}Weight based on the name and
// callsite to compute node and edge weights which will be used later on to
// create edges for WeightedCG.
//
// Caller should ignore the profile if p.TotalNodeWeight == 0 || p.TotalEdgeWeight == 0.
func (p *Profile) processprofileGraph(g *graph.Graph) error {
	nFlat := make(map[string]int64)
	nCum := make(map[string]int64)
	seenStartLine := false

	// Accummulate weights for the same node.
	for _, n := range g.Nodes {
		canonicalName := n.Info.Name
		nFlat[canonicalName] += n.FlatValue()
		nCum[canonicalName] += n.CumValue()
	}

	// Process graph and build various node and edge maps which will
	// be consumed by AST walk.
	for _, n := range g.Nodes {
		seenStartLine = seenStartLine || n.Info.StartLine != 0

		p.TotalNodeWeight += n.FlatValue()
		canonicalName := n.Info.Name
		// Create the key to the nodeMapKey.
		nodeinfo := NodeMapKey{
			CallerName:     canonicalName,
			CallSiteOffset: n.Info.Lineno - n.Info.StartLine,
		}

		for _, e := range n.Out {
			p.TotalEdgeWeight += e.WeightValue()
			nodeinfo.CalleeName = e.Dest.Info.Name
			if w, ok := p.NodeMap[nodeinfo]; ok {
				w.EWeight += e.WeightValue()
			} else {
				weights := new(Weights)
				weights.NFlat = nFlat[canonicalName]
				weights.NCum = nCum[canonicalName]
				weights.EWeight = e.WeightValue()
				p.NodeMap[nodeinfo] = weights
			}
		}
	}

	if p.TotalNodeWeight == 0 || p.TotalEdgeWeight == 0 {
		return nil // accept but ignore profile with no samples.
	}

	if !seenStartLine {
		// TODO(prattmic): If Function.start_line is missing we could
		// fall back to using absolute line numbers, which is better
		// than nothing.
		return fmt.Errorf("profile missing Function.start_line data (Go version of profiled application too old? Go 1.20+ automatically adds this to profiles)")
	}

	return nil
}

// initializeIRGraph builds the IRGraph by visiting all the ir.Func in decl list
// of a package.
func (p *Profile) initializeIRGraph() {
	// Bottomup walk over the function to create IRGraph.
	ir.VisitFuncsBottomUp(typecheck.Target.Decls, func(list []*ir.Func, recursive bool) {
		for _, fn := range list {
			p.VisitIR(fn)
		}
	})

	// Add additional edges for indirect calls. This must be done second so
	// that IRNodes is fully populated (see the dummy node TODO in
	// addIndirectEdges).
	//
	// TODO(prattmic): VisitIR above populates the graph via direct calls
	// discovered via the IR. addIndirectEdges populates the graph via
	// calls discovered via the profile. This combination of opposite
	// approaches is a bit awkward, particularly because direct calls are
	// discoverable via the profile as well. Unify these into a single
	// approach.
	p.addIndirectEdges()
}

// VisitIR traverses the body of each ir.Func and use NodeMap to determine if
// we need to add an edge from ir.Func and any node in the ir.Func body.
func (p *Profile) VisitIR(fn *ir.Func) {
	g := p.WeightedCG

	if g.IRNodes == nil {
		g.IRNodes = make(map[string]*IRNode)
	}

	name := ir.LinkFuncName(fn)
	node, ok := g.IRNodes[name]
	if !ok {
		node = &IRNode{
			AST: fn,
		}
		g.IRNodes[name] = node
	}

	// Recursively walk over the body of the function to create IRGraph edges.
	p.createIRGraphEdge(fn, node, name)
}

// NodeLineOffset returns the line offset of n in fn.
func NodeLineOffset(n ir.Node, fn *ir.Func) int {
	// See "A note on line numbers" at the top of the file.
	line := int(base.Ctxt.InnermostPos(n.Pos()).RelLine())
	startLine := int(base.Ctxt.InnermostPos(fn.Pos()).RelLine())
	return line - startLine
}

// addIREdge adds an edge between caller and new node that points to `callee`
// based on the profile-graph and NodeMap.
func (p *Profile) addIREdge(callerNode *IRNode, callerName string, call ir.Node, callee *ir.Func) {
	g := p.WeightedCG

	calleeName := ir.LinkFuncName(callee)
	calleeNode, ok := g.IRNodes[calleeName]
	if !ok {
		calleeNode = &IRNode{
			AST: callee,
		}
		g.IRNodes[calleeName] = calleeNode
	}

	nodeinfo := NodeMapKey{
		CallerName:     callerName,
		CalleeName:     calleeName,
		CallSiteOffset: NodeLineOffset(call, callerNode.AST),
	}

	var weight int64
	if weights, ok := p.NodeMap[nodeinfo]; ok {
		weight = weights.EWeight
	}

	// Add edge in the IRGraph from caller to callee.
	edge := &IREdge{
		Src:            callerNode,
		Dst:            calleeNode,
		Weight:         weight,
		CallSiteOffset: nodeinfo.CallSiteOffset,
	}

	if callerNode.OutEdges == nil {
		callerNode.OutEdges = make(map[NodeMapKey]*IREdge)
	}
	callerNode.OutEdges[nodeinfo] = edge
}

// addIndirectEdges adds indirect call edges found in the profile to the graph,
// to be used for devirtualization.
//
// targetDeclFuncs is the set of functions in typecheck.Target.Decls. Only
// edges from these functions will be added.
//
// Devirtualization is only applied to typecheck.Target.Decls functions, so there
// is no need to add edges from other functions.
//
// N.B. despite the name, addIndirectEdges will add any edges discovered via
// the profile. We don't know for sure that they are indirect, but assume they
// are since direct calls would already be added. (e.g., direct calls that have
// been deleted from source since the profile was taken would be added here).
//
// TODO(prattmic): Devirtualization runs before inlining, so we can't devirtualize
// calls inside inlined call bodies. If we did add that, we'd need edges from
// inlined bodies as well.
func (p *Profile) addIndirectEdges() {
	g := p.WeightedCG

	// g.IRNodes is populated with the set of functions in the local
	// package build by VisitIR. We want to filter for local functions
	// below, but we also add unknown callees to IRNodes as we go. So make
	// an initial copy of IRNodes to recall just the local functions.
	localNodes := make(map[string]*IRNode, len(g.IRNodes))
	for k, v := range g.IRNodes {
		localNodes[k] = v
	}

	for key, weights := range p.NodeMap {
		// All callers in the local package build were added to IRNodes
		// in VisitIR. If a caller isn't in the local package build we
		// can skip adding edges, since we won't be devirtualizing in
		// them anyway. This keeps the graph smaller.
		callerNode, ok := localNodes[key.CallerName]
		if !ok {
			continue
		}

		// Already handled this edge?
		if _, ok := callerNode.OutEdges[key]; ok {
			continue
		}

		calleeNode, ok := g.IRNodes[key.CalleeName]
		if !ok {
			// IR is missing for this callee. Most likely this is
			// because the callee isn't in the transitive deps of
			// this package.
			//
			// Record this call anyway. If this is the hottest,
			// then we want to skip devirtualization rather than
			// devirtualizing to the second most common callee.
			//
			// TODO(prattmic): VisitIR populates IRNodes with all
			// of the functions discovered via local package
			// function declarations and calls. Thus we could miss
			// functions that are available in export data of
			// transitive deps, but aren't directly reachable. We
			// need to do a lookup directly from package export
			// data to get complete coverage.
			calleeNode = &IRNode{
				LinkerSymbolName: key.CalleeName,
				// TODO: weights? We don't need them.
			}
			// Add dummy node back to IRNodes. We don't need this
			// directly, but PrintWeightedCallGraphDOT uses these
			// to print nodes.
			g.IRNodes[key.CalleeName] = calleeNode
		}
		edge := &IREdge{
			Src:            callerNode,
			Dst:            calleeNode,
			Weight:         weights.EWeight,
			CallSiteOffset: key.CallSiteOffset,
		}

		if callerNode.OutEdges == nil {
			callerNode.OutEdges = make(map[NodeMapKey]*IREdge)
		}
		callerNode.OutEdges[key] = edge
	}
}

// createIRGraphEdge traverses the nodes in the body of ir.Func and adds edges
// between the callernode which points to the ir.Func and the nodes in the
// body.
func (p *Profile) createIRGraphEdge(fn *ir.Func, callernode *IRNode, name string) {
	ir.VisitList(fn.Body, func(n ir.Node) {
		switch n.Op() {
		case ir.OCALLFUNC:
			call := n.(*ir.CallExpr)
			// Find the callee function from the call site and add the edge.
			callee := DirectCallee(call.X)
			if callee != nil {
				p.addIREdge(callernode, name, n, callee)
			}
		case ir.OCALLMETH:
			call := n.(*ir.CallExpr)
			// Find the callee method from the call site and add the edge.
			callee := ir.MethodExprName(call.X).Func
			p.addIREdge(callernode, name, n, callee)
		}
	})
}

// WeightInPercentage converts profile weights to a percentage.
func WeightInPercentage(value int64, total int64) float64 {
	return (float64(value) / float64(total)) * 100
}

// PrintWeightedCallGraphDOT prints IRGraph in DOT format.
func (p *Profile) PrintWeightedCallGraphDOT(edgeThreshold float64) {
	fmt.Printf("\ndigraph G {\n")
	fmt.Printf("forcelabels=true;\n")

	// List of functions in this package.
	funcs := make(map[string]struct{})
	ir.VisitFuncsBottomUp(typecheck.Target.Decls, func(list []*ir.Func, recursive bool) {
		for _, f := range list {
			name := ir.LinkFuncName(f)
			funcs[name] = struct{}{}
		}
	})

	// Determine nodes of DOT.
	//
	// Note that ir.Func may be nil for functions not visible from this
	// package.
	nodes := make(map[string]*ir.Func)
	for name := range funcs {
		if n, ok := p.WeightedCG.IRNodes[name]; ok {
			for _, e := range n.OutEdges {
				if _, ok := nodes[e.Src.Name()]; !ok {
					nodes[e.Src.Name()] = e.Src.AST
				}
				if _, ok := nodes[e.Dst.Name()]; !ok {
					nodes[e.Dst.Name()] = e.Dst.AST
				}
			}
			if _, ok := nodes[n.Name()]; !ok {
				nodes[n.Name()] = n.AST
			}
		}
	}

	// Print nodes.
	for name, ast := range nodes {
		if _, ok := p.WeightedCG.IRNodes[name]; ok {
			style := "solid"
			if ast == nil {
				style = "dashed"
			}

			if ast != nil && ast.Inl != nil {
				fmt.Printf("\"%v\" [color=black, style=%s, label=\"%v,inl_cost=%d\"];\n", name, style, name, ast.Inl.Cost)
			} else {
				fmt.Printf("\"%v\" [color=black, style=%s, label=\"%v\"];\n", name, style, name)
			}
		}
	}
	// Print edges.
	ir.VisitFuncsBottomUp(typecheck.Target.Decls, func(list []*ir.Func, recursive bool) {
		for _, f := range list {
			name := ir.LinkFuncName(f)
			if n, ok := p.WeightedCG.IRNodes[name]; ok {
				for _, e := range n.OutEdges {
					style := "solid"
					if e.Dst.AST == nil {
						style = "dashed"
					}
					color := "black"
					edgepercent := WeightInPercentage(e.Weight, p.TotalEdgeWeight)
					if edgepercent > edgeThreshold {
						color = "red"
					}

					fmt.Printf("edge [color=%s, style=%s];\n", color, style)
					fmt.Printf("\"%v\" -> \"%v\" [label=\"%.2f\"];\n", n.Name(), e.Dst.Name(), edgepercent)
				}
			}
		}
	})
	fmt.Printf("}\n")
}

// DirectCallee takes a function-typed expression and returns the underlying
// function that it refers to if statically known. Otherwise, it returns nil.
//
// Equivalent to inline.inlCallee without calling CanInline on closures.
func DirectCallee(fn ir.Node) *ir.Func {
	fn = ir.StaticValue(fn)
	switch fn.Op() {
	case ir.OMETHEXPR:
		fn := fn.(*ir.SelectorExpr)
		n := ir.MethodExprName(fn)
		// Check that receiver type matches fn.X.
		// TODO(mdempsky): Handle implicit dereference
		// of pointer receiver argument?
		if n == nil || !types.Identical(n.Type().Recv().Type, fn.X.Type()) {
			return nil
		}
		return n.Func
	case ir.ONAME:
		fn := fn.(*ir.Name)
		if fn.Class == ir.PFUNC {
			return fn.Func
		}
	case ir.OCLOSURE:
		fn := fn.(*ir.ClosureExpr)
		c := fn.Func
		return c
	}
	return nil
}