Revert "Revert "compiler-wrapper: adds an IWYU component""

This reverts commit 537f2ecfe7bc9eb996cd4abedae26bf699ab54dd.

It also fixes the logic bug which caused the reert in the first place.
IWYU flag prcessing was removing the file names if IWYU was not needed,
which meant clang tidy calls were not getting the file names added to
the back of the command. It also fixes the header to the most recent
style guidance and updates some golang variable names to make sure we
pass go lint.

BUG=b:237320348
TEST=go test

Change-Id: I14885fb90f97d2fb483ae1a01e14c44050852dc3
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/3997486
Commit-Queue: Ryan Beltran <ryanbeltran@chromium.org>
Reviewed-by: Ryan Beltran <ryanbeltran@chromium.org>
Reviewed-by: Manoj Gupta <manojgupta@chromium.org>
Tested-by: Ryan Beltran <ryanbeltran@chromium.org>
diff --git a/compiler_wrapper/compiler_wrapper.go b/compiler_wrapper/compiler_wrapper.go
index 1386374..dcaada9 100644
--- a/compiler_wrapper/compiler_wrapper.go
+++ b/compiler_wrapper/compiler_wrapper.go
@@ -151,6 +151,7 @@
 		}
 	} else {
 		cSrcFile, tidyFlags, tidyMode := processClangTidyFlags(mainBuilder)
+		cSrcFile, iwyuFlags, iwyuMode := processIWYUFlags(mainBuilder)
 		if mainBuilder.target.compilerType == clangType {
 			err := prepareClangCommand(mainBuilder)
 			if err != nil {
@@ -176,6 +177,20 @@
 					return 0, err
 				}
 			}
+
+			if iwyuMode != iwyuModeNone {
+				if iwyuMode == iwyuModeError {
+					panic(fmt.Sprintf("Unknown IWYU mode"))
+				}
+
+				allowCCache = false
+				clangCmdWithoutRemoteBuildAndCCache := mainBuilder.build()
+				err := runIWYU(env, clangCmdWithoutRemoteBuildAndCCache, cSrcFile, iwyuFlags)
+				if err != nil {
+					return 0, err
+				}
+			}
+
 			if remoteBuildUsed, err = processRemoteBuildAndCCacheFlags(allowCCache, mainBuilder); err != nil {
 				return 0, err
 			}
diff --git a/compiler_wrapper/iwyu_flag.go b/compiler_wrapper/iwyu_flag.go
new file mode 100644
index 0000000..d13d114
--- /dev/null
+++ b/compiler_wrapper/iwyu_flag.go
@@ -0,0 +1,156 @@
+// Copyright 2022 The ChromiumOS Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"path/filepath"
+	"strings"
+)
+
+type useIWYUMode int
+
+const iwyuCrashSubstring = "PLEASE submit a bug report"
+
+const (
+	iwyuModeNone useIWYUMode = iota
+	iwyuModeAll
+	iwyuModeError
+)
+
+var srcFileSuffixes = []string{
+	".c",
+	".cc",
+	".cpp",
+	".C",
+	".cxx",
+	".c++",
+}
+
+func findWithIWYUFlag(args []builderArg) (string, []builderArg) {
+	for i := range args {
+		if args[i].value == "--with-iwyu" {
+			args = append(args[:i], args[i+1:]...)
+			return "1", args
+		}
+	}
+	return "", args
+}
+
+func processIWYUFlags(builder *commandBuilder) (cSrcFile string, iwyuFlags []string, mode useIWYUMode) {
+	builder.transformArgs(func(arg builderArg) string {
+		const prefix = "-iwyu-flag="
+		if !strings.HasPrefix(arg.value, prefix) {
+			return arg.value
+		}
+
+		iwyuFlags = append(iwyuFlags, arg.value[len(prefix):])
+		return ""
+	})
+
+	cSrcFile = ""
+	lastArg := ""
+	for _, arg := range builder.args {
+		if lastArg != "-o" {
+			for _, suffix := range srcFileSuffixes {
+				if strings.HasSuffix(arg.value, suffix) {
+					cSrcFile = arg.value
+					break
+				}
+			}
+		}
+		lastArg = arg.value
+	}
+
+	if cSrcFile == "" {
+		return "", iwyuFlags, iwyuModeNone
+	}
+
+	withIWYU, _ := builder.env.getenv("WITH_IWYU")
+	if withIWYU == "" {
+		withIWYU, builder.args = findWithIWYUFlag(builder.args)
+		if withIWYU == "" {
+			return cSrcFile, iwyuFlags, iwyuModeNone
+		}
+	}
+
+	if withIWYU != "1" {
+		return cSrcFile, iwyuFlags, iwyuModeError
+	}
+
+	return cSrcFile, iwyuFlags, iwyuModeAll
+}
+
+func calcIWYUInvocation(env env, clangCmd *command, cSrcFile string, iwyuFlags ...string) (*command, error) {
+	resourceDir, err := getClangResourceDir(env, clangCmd.Path)
+	if err != nil {
+		return nil, err
+	}
+
+	iwyuPath := filepath.Join(filepath.Dir(clangCmd.Path), "include-what-you-use")
+	args := append([]string{}, iwyuFlags...)
+	args = append(args, "-resource-dir="+resourceDir)
+	args = append(args, clangCmd.Args...)
+
+	for i := 0; i < len(args); i++ {
+		for j := 0; j < len(srcFileSuffixes); j++ {
+			if strings.HasSuffix(args[i], srcFileSuffixes[j]) {
+				args = append(args[:i], args[i+1:]...)
+				break
+			}
+		}
+	}
+	args = append(args, cSrcFile)
+
+	return &command{
+		Path:       iwyuPath,
+		Args:       args,
+		EnvUpdates: clangCmd.EnvUpdates,
+	}, nil
+}
+
+func runIWYU(env env, clangCmd *command, cSrcFile string, extraIWYUFlags []string) error {
+	extraIWYUFlags = append(extraIWYUFlags, "-Xiwyu", "--mapping_file=/usr/share/include-what-you-use/libcxx.imp", "-Xiwyu", "--no_fwd_decls")
+	iwyuCmd, err := calcIWYUInvocation(env, clangCmd, cSrcFile, extraIWYUFlags...)
+	if err != nil {
+		return fmt.Errorf("calculating include-what-you-use invocation: %v", err)
+	}
+
+	// Note: We pass nil as stdin as we checked before that the compiler
+	// was invoked with a source file argument.
+	var stderr bytes.Buffer
+	stderrWriter := bufio.NewWriter(&stderr)
+	exitCode, err := wrapSubprocessErrorWithSourceLoc(iwyuCmd,
+		env.run(iwyuCmd, nil, nil, stderrWriter))
+	stderrMessage := stderr.String()
+	fmt.Fprintln(env.stderr(), stderrMessage)
+
+	if err == nil && exitCode != 0 {
+		// Note: We continue on purpose when include-what-you-use fails
+		// to maintain compatibility with the previous wrapper.
+		fmt.Fprintln(env.stderr(), "include-what-you-use failed")
+	}
+
+	var path strings.Builder
+	path.WriteString(strings.TrimSuffix(iwyuCmd.Path, "include-what-you-use"))
+	path.WriteString("fix_includes.py")
+	fixIncludesCmd := &command{
+		Path:       path.String(),
+		Args:       []string{"--nocomment"},
+		EnvUpdates: clangCmd.EnvUpdates,
+	}
+
+	exitCode, err = wrapSubprocessErrorWithSourceLoc(fixIncludesCmd,
+		env.run(fixIncludesCmd, strings.NewReader(stderrMessage), env.stdout(), env.stderr()))
+	if err == nil && exitCode != 0 {
+		// Note: We continue on purpose when include-what-you-use fails
+		// to maintain compatibility with the previous wrapper.
+		fmt.Fprint(env.stderr(), "include-what-you-use failed")
+	}
+
+	return err
+}
diff --git a/compiler_wrapper/iwyu_flag_test.go b/compiler_wrapper/iwyu_flag_test.go
new file mode 100644
index 0000000..7613594
--- /dev/null
+++ b/compiler_wrapper/iwyu_flag_test.go
@@ -0,0 +1,135 @@
+// Copyright 2022 The ChromiumOS Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package main
+
+import (
+	"errors"
+	"io"
+	"strings"
+	"testing"
+)
+
+func TestIWYUArgOrder(t *testing.T) {
+	withIWYUTestContext(t, func(ctx *testContext) {
+		ctx.cmdMock = func(cmd *command, stdin io.Reader, stdout io.Writer, stderr io.Writer) error {
+			if ctx.cmdCount == 2 {
+				if err := verifyArgOrder(cmd, "-checks=.*", mainCc, "--", "-resource-dir=.*", mainCc, "--some_arg"); err != nil {
+					return err
+				}
+			}
+			return nil
+		}
+		ctx.must(callCompiler(ctx, ctx.cfg,
+			ctx.newCommand(clangX86_64, mainCc, "--some_arg")))
+		if ctx.cmdCount < 2 {
+			t.Error("expected multiple calls.")
+		}
+	})
+}
+
+func TestIgnoreNonZeroExitCodeFromIWYU(t *testing.T) {
+	withIWYUTestContext(t, func(ctx *testContext) {
+		ctx.cmdMock = func(cmd *command, stdin io.Reader, stdout io.Writer, stderr io.Writer) error {
+			if ctx.cmdCount == 2 {
+				return newExitCodeError(23)
+			}
+			return nil
+		}
+		ctx.must(callCompiler(ctx, ctx.cfg,
+			ctx.newCommand(clangX86_64, mainCc)))
+		stderr := ctx.stderrString()
+		if err := verifyNonInternalError(stderr, "include-what-you-use failed"); err != nil {
+			t.Error(err)
+		}
+	})
+}
+
+func TestReportGeneralErrorsFromIWYU(t *testing.T) {
+	withIWYUTestContext(t, func(ctx *testContext) {
+		ctx.cmdMock = func(cmd *command, stdin io.Reader, stdout io.Writer, stderr io.Writer) error {
+			if ctx.cmdCount > 1 {
+				return errors.New("someerror")
+			}
+			return nil
+		}
+		stderr := ctx.mustFail(callCompiler(ctx, ctx.cfg,
+			ctx.newCommand(clangX86_64, mainCc)))
+		if err := verifyInternalError(stderr); err != nil {
+			t.Fatal(err)
+		}
+		if !strings.Contains(stderr, "someerror") {
+			t.Errorf("unexpected error. Got: %s", stderr)
+		}
+	})
+}
+
+func TestUseIWYUBasedOnFileExtension(t *testing.T) {
+	withIWYUTestContext(t, func(ctx *testContext) {
+		testData := []struct {
+			args []string
+			iwyu bool
+		}{
+			{[]string{"main.cc"}, true},
+			{[]string{"main.cc"}, true},
+			{[]string{"main.C"}, true},
+			{[]string{"main.cxx"}, true},
+			{[]string{"main.c++"}, true},
+			{[]string{"main.xy"}, false},
+			{[]string{"-o", "main.cc"}, false},
+			{[]string{}, false},
+		}
+		for _, tt := range testData {
+			ctx.cmdCount = 0
+			ctx.must(callCompiler(ctx, ctx.cfg,
+				ctx.newCommand(clangX86_64, tt.args...)))
+			if ctx.cmdCount > 1 && !tt.iwyu {
+				t.Errorf("expected a call to iwyu but got none for args %s", tt.args)
+			}
+			if ctx.cmdCount == 1 && tt.iwyu {
+				t.Errorf("expected no call to iwyu but got one for args %s", tt.args)
+			}
+		}
+	})
+}
+
+func TestIWYUFiltersIWYUFlags(t *testing.T) {
+	withIWYUTestContext(t, func(ctx *testContext) {
+		addedFlag := "--some_iwyu_flag=flag"
+		ctx.cmdMock = func(cmd *command, stdin io.Reader, stdout io.Writer, stderr io.Writer) error {
+			switch ctx.cmdCount {
+			case 1:
+				if err := verifyPath(cmd, "usr/bin/clang"); err != nil {
+					t.Error(err)
+				} else if err := verifyArgCount(cmd, 0, addedFlag); err != nil {
+					t.Error(err)
+				}
+				return nil
+			case 2:
+				if err := verifyPath(cmd, "usr/bin/include-what-you-use"); err != nil {
+					t.Error(err)
+				} else if verifyArgCount(cmd, 1, addedFlag); err != nil {
+					t.Error(err)
+				}
+				return nil
+			default:
+				return nil
+			}
+		}
+		cmd := ctx.must(callCompiler(ctx, ctx.cfg, ctx.newCommand(clangX86_64, mainCc, "-iwyu-flag="+addedFlag)))
+		if ctx.cmdCount < 2 {
+			t.Errorf("expected multiple calls.")
+		}
+		if err := verifyPath(cmd, "usr/bin/clang"); err != nil {
+			t.Error(err)
+		}
+	})
+}
+
+func withIWYUTestContext(t *testing.T, work func(ctx *testContext)) {
+	withTestContext(t, func(ctx *testContext) {
+		ctx.env = []string{"WITH_IWYU=1"}
+		work(ctx)
+	})
+}