go/src/infra/cmd/package_index/main.go

// Copyright 2020 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"cmp"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"text/tabwriter"
	"time"
	"unique"

	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/logging/gologger"

	kpb "go.chromium.org/infra/cmd/package_index/kythe/proto"
)

// Size of buffered channels.
var chanSize = 1000

// Number of goroutines to use in parallel processing.
var numRoutines = 32

var (
	projectFlag         = flag.String("project", "chromium", "Project this kzip should be generated for. This should be set to 'chromium', 'chrome', or 'chromiumos'. Defaults to 'chromium'.")
	outputFlag          = flag.String("path_to_archive_output", "", "Path to index pack archive to be generated.")
	compDBFlag          = flag.String("path_to_compdb", "", "Path to the compilation database.")
	gnFlag              = flag.String("path_to_gn_targets", "", "Path to the gn targets json file.")
	corpusFlag          = flag.String("corpus", "", "Kythe corpus to use for the vname.")
	existingKzipsFlag   = flag.String("path_to_java_kzips", "", "Path to already generated java kzips which will be included in the final index pack.")
	buildFlag           = flag.String("build_config", "", "Build config to use in the unit file.")
	clangTargetArchFlag = flag.String("clang_target_arch", "", "Target architecture to provide via -target in any clang commands.")
	checkoutFlag        = flag.String("checkout_dir", "", "Root of the repository.")
	outDirFlag          = flag.String("out_dir", "src/out/Debug", "Output directory from which compilation is run.")
	filepathsFlag       = flag.Bool("keep_filepaths_files", false, "Keep the .filepaths files used for index pack generation.")
	verboseFlag         = flag.Bool("verbose", false, "Print the details of every file being written to the index pack.")
	useSisoFlag         = flag.Bool("use_siso", false, "Call siso query deps to get clang files instead of reading from *.filepaths")

	allowedProjects = stringset.NewFromSlice(
		"chromium",
		"chromiumos",
		"chrome",
		"chromeos")
)

// This should be same as
// https://chromium.googlesource.com/build/+/siso/v1.4.22/siso/subcmd/query/deps.go#77
// but we want to avoid depending on build repo from here.
type depInfo struct {
	Target string   `json:"target"`
	Deps   []string `json:"deps"`
}

// Only built when use_siso = true.
// Key: source file, value: list of dependencies.
type targetDepsMap = map[unique.Handle[string]][]unique.Handle[string]

// validateFlags checks that the required flags are present.
func validateFlags(ctx context.Context) {
	flagErr := false

	if !allowedProjects.Has(*projectFlag) {
		logging.Errorf(ctx, "project not supported. [%s] are the only supported projects.", strings.Join(allowedProjects.ToSortedSlice(), ", "))
		flagErr = true
	}

	if *outputFlag == "" {
		logging.Errorf(ctx, "path_to_archive_output flag required.")
		flagErr = true
	}

	if *compDBFlag == "" {
		logging.Errorf(ctx, "path_to_compdb flag required.")
		flagErr = true
	}

	if *gnFlag == "" {
		logging.Errorf(ctx, "path_to_gn_targets flag required.")
		flagErr = true
	}

	if *corpusFlag == "" {
		logging.Errorf(ctx, "corpus flag required.")
		flagErr = true
	}

	if *checkoutFlag == "" {
		logging.Errorf(ctx, "checkout_dir flag required.")
		flagErr = true
	}

	if flagErr {
		panic("Flags are missing or not supported")
	}
}

// Return if a filename is a source file with certain file extensions.
func isSourceFile(filename string) bool {
	ext := filepath.Ext(filename)
	switch ext {
	case ".c", ".cc", ".cxx", ".cpp", ".m", ".mm", ".S":
		return true
	default:
		return false
	}
}

func loadSisoJsonDeps(depsReader io.Reader) (targetDepsMap, error) {
	m := make(targetDepsMap)

	var target depInfo
	decoder := json.NewDecoder(depsReader)
	for {
		// Clean up the target variable to reuse.
		target.Target = ""
		target.Deps = target.Deps[:0]

		err := decoder.Decode(&target)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return m, fmt.Errorf("fail to decode record with error: %v", err)
		}
		// Only process obj/ targets.
		if len(target.Deps) == 0 || !strings.HasPrefix(target.Target, "obj/") {
			continue
		}
		var sourceFileHandle unique.Handle[string]
		targetHandles := make([]unique.Handle[string], 0, len(target.Deps))
		for _, dep := range target.Deps {
			h := unique.Make(dep)
			// Find the first source file.
			// It is usually the first/second element of the dependency list.
			if sourceFileHandle == (unique.Handle[string]{}) && isSourceFile(dep) {
				sourceFileHandle = h
			}
			targetHandles = append(targetHandles, h)
		}

		if sourceFileHandle == (unique.Handle[string]{}) {
			// If no source file, just ignore the current group.
			continue
		}
		value, ok := m[sourceFileHandle]
		if !ok {
			m[sourceFileHandle] = targetHandles
			continue
		}
		// Merge the dependencies if source file is previously encountered.
		set := make(map[unique.Handle[string]]struct{})
		for _, val := range value {
			set[val] = struct{}{}
		}
		for _, val := range targetHandles {
			set[val] = struct{}{}
		}
		result := make([]unique.Handle[string], 0, len(set))
		for key := range set {
			result = append(result, key)
		}
		slices.SortFunc(result, func(a, b unique.Handle[string]) int {
			return cmp.Compare(a.Value(), b.Value())
		})
		m[sourceFileHandle] = result
	}

	return m, nil
}

func main() {
	ctx := gologger.StdConfig.Use(context.Background())
	flag.Parse()
	validateFlags(ctx)

	// Remove the old zip archive (if it exists). This avoids the new index
	// pack being added to the old zip archive.
	if _, err := os.Stat(*outputFlag); err == nil { // Old kzip exists.
		err = os.Remove(*outputFlag)
		if err != nil {
			panic(err)
		}
	}

	// Setup.
	if *verboseFlag {
		ctx = logging.SetLevel(ctx, logging.Debug)
	}
	logging.Infof(ctx, "%s: Index generation...", time.Now().Format("15:04:05"))
	rootPath, err := filepath.Abs(filepath.Join(*checkoutFlag, ".."))
	if err != nil {
		panic(err)
	}
	ip := newIndexPack(ctx, *outputFlag, rootPath, *outDirFlag, *compDBFlag,
		*gnFlag, *existingKzipsFlag, *corpusFlag, *buildFlag, *clangTargetArchFlag)

	// Process existing kzips.
	existingKzipChannel := make(chan string, chanSize)
	go func() {
		if ip.existingJavaKzipsPath != "" {
			err := ip.mergeExistingKzips(existingKzipChannel)
			if err != nil {
				panic(err)
			}
		} else {
			close(existingKzipChannel)
		}
	}()

	var kzipEntryWg sync.WaitGroup
	kzipEntryChannel := make(chan kzipEntry, 100) // Channel size is reduced for chromiumos builder.
	kzipSet := NewConcurrentSet(0)
	kzipEntryWg.Add(1)
	go func() {
		err := ip.processExistingKzips(ctx, existingKzipChannel, kzipEntryChannel, kzipSet)
		if err != nil {
			panic(err)
		}
		kzipEntryWg.Done()
	}()

	// Process targets.
	unitProtoChannel := make(chan *kpb.CompilationUnit, chanSize)
	dataFileChannel := make(chan string, chanSize)

	// Process clang targets.
	clangTargets := NewClangTargets(ip.compDBPath)
	clangTargets.DataWg.Add(numRoutines)
	clangTargets.KzipDataWg.Add(numRoutines)
	clangTargets.UnitWg.Add(numRoutines)

	// Load siso output into memory when enabled.
	var targetDepsMap targetDepsMap
	if *useSisoFlag {
		targetDepsMap, err = loadSisoJsonDeps(os.Stdin)
		if err != nil {
			panic(err)
		}
	}

	// Process clang targets.
	for range numRoutines {
		go func() {
			err := clangTargets.ProcessClangTargets(ip.ctx, ip.rootPath, ip.outDir, ip.corpus,
				ip.buildConfig, ip.clangTargetArch, ip.hashMaps, dataFileChannel, unitProtoChannel, *useSisoFlag, targetDepsMap)
			if err != nil {
				panic(err)
			}
		}()
	}

	// Process GN targets.
	gnTargets := NewGnTargets(ip.gnTargetsPath)
	gnTargets.DataWg.Add(numRoutines)
	gnTargets.KzipDataWg.Add(numRoutines)
	gnTargets.UnitWg.Add(numRoutines)
	for range numRoutines {
		go func() {
			err := gnTargets.ProcessGnTargets(ip.ctx, ip.rootPath, ip.outDir, ip.corpus, ip.buildConfig, ip.hashMaps,
				dataFileChannel, unitProtoChannel)
			if err != nil {
				panic(err)
			}
		}()
	}

	// Convert data files to kzipEntries.
	kzipEntryWg.Add(numRoutines)
	for range numRoutines {
		go func() {
			if err := ip.dataFileToKzipEntry(ctx, dataFileChannel, kzipEntryChannel); err != nil {
				panic(err)
			}
			kzipEntryWg.Done()

			// Signal targets to start unit proto processing.
			gnTargets.KzipDataWg.Done()
			clangTargets.KzipDataWg.Done()
		}()
	}

	// Convert unit protos to kzipEntries.
	kzipEntryWg.Add(numRoutines)
	for range numRoutines {
		go func() {
			ip.unitFileToKzipEntry(ctx, unitProtoChannel, kzipEntryChannel)
			kzipEntryWg.Done()
		}()
	}

	// Close dataFileChannel and unitProtoChannel after all targets have been processed and sent.
	go func() {
		gnTargets.DataWg.Wait()
		clangTargets.DataWg.Wait()
		close(dataFileChannel)
	}()
	go func() {
		gnTargets.UnitWg.Wait()
		clangTargets.UnitWg.Wait()
		close(unitProtoChannel)
	}()

	// Close kzipEntryChannel after all kzip entries have been sent.
	go func() {
		kzipEntryWg.Wait()
		close(kzipEntryChannel)
	}()

	// Write all data file and unit proto entries to kzip.
	err = ip.writeToKzip(kzipEntryChannel)

	// Print stats regardless of success.
	tw := tabwriter.NewWriter(os.Stdout, 0, 0, 3, '-', tabwriter.AlignRight|tabwriter.Debug)
	fmt.Fprintln(tw, "Language\tCompilationUnits\tRequiredInputs\tSourceFiles")
	numUnits := 0
	numRequiredInputs := 0
	numSourceFiles := 0
	for l, st := range ip.stats {
		nri := len(st.requiredInputs)
		nsf := len(st.sourceFiles)
		fmt.Fprintf(tw, "%s\t%d\t%d\t%d\n", l, st.numCompilationUnits, nri, nsf)
		numUnits += st.numCompilationUnits
		numRequiredInputs += nri
		numSourceFiles += nsf
	}
	fmt.Fprintf(tw, "Total\t%d\t%d\t%d\n", numUnits, numRequiredInputs, numSourceFiles)
	if tw.Flush() != nil {
		logging.Warningf(ctx, "failed to flush stats to stdout")
	}
	if err != nil {
		panic(err)
	}

	// Clean up.
	if !*filepathsFlag {
		// Remove all *.filepaths files.
		removeFilepathsFiles(ip.ctx, filepath.Join(rootPath, "src"))
	}
	logging.Infof(ctx, "%s: Done.", time.Now().Format("15:04:05"))
	os.Exit(0)
}