Revamp flag handling in pkg/prog.

Instead of putting all possible flags in prog.Flags, flags are now registered by
the individual subprograms. The 3 flags -sock, -db and -json are shared by
multiple subprograms and still handled by the prog package.

This new design allows separating the support for -cpuprofile into a separate
subprogram, which is no longer included by the default entry point, making the
binary slightly smaller. A new entrypoint "withpprof" is created.

Also include the LSP subprogram in the nodaemon entry point.

cmd/elvish/main.go

@@ -18,6 +18,6 @@ func main() {
 	os.Exit(prog.Run(
 		[3]*os.File{os.Stdin, os.Stdout, os.Stderr}, os.Args,
 		prog.Composite(
-			buildinfo.Program, daemon.Program, lsp.Program,
-			shell.Program{ActivateDaemon: daemon.Activate})))
+			&buildinfo.Program{}, &daemon.Program{}, &lsp.Program{},
+			&shell.Program{ActivateDaemon: daemon.Activate})))
 }

cmd/nodaemon/elvish/main.go

@@ -3,10 +3,10 @@
 package main
 
 import (
-	"errors"
 	"os"
 
 	"src.elv.sh/pkg/buildinfo"
+	"src.elv.sh/pkg/lsp"
 	"src.elv.sh/pkg/prog"
 	"src.elv.sh/pkg/shell"
 )
@@ -14,16 +14,5 @@ import (
 func main() {
 	os.Exit(prog.Run(
 		[3]*os.File{os.Stdin, os.Stdout, os.Stderr}, os.Args,
-		prog.Composite(buildinfo.Program, daemonStub{}, shell.Program{})))
-}
-
-var errNoDaemon = errors.New("daemon is not supported in this build")
-
-type daemonStub struct{}
-
-func (daemonStub) Run(fds [3]*os.File, f *prog.Flags, args []string) error {
-	if f.Daemon {
-		return errNoDaemon
-	}
-	return prog.ErrNotSuitable
+		prog.Composite(&buildinfo.Program{}, &lsp.Program{}, &shell.Program{})))
 }

cmd/withpprof/elvish/main.go

@@ -0,0 +1,22 @@
+// Command elvish is an alternative main program of Elvish that supports writing
+// pprof profiles.
+package main
+
+import (
+	"os"
+
+	"src.elv.sh/pkg/buildinfo"
+	"src.elv.sh/pkg/daemon"
+	"src.elv.sh/pkg/lsp"
+	"src.elv.sh/pkg/pprof"
+	"src.elv.sh/pkg/prog"
+	"src.elv.sh/pkg/shell"
+)
+
+func main() {
+	os.Exit(prog.Run(
+		[3]*os.File{os.Stdin, os.Stdout, os.Stderr}, os.Args,
+		prog.Composite(
+			&pprof.Program{}, &buildinfo.Program{}, &daemon.Program{}, &lsp.Program{},
+			&shell.Program{ActivateDaemon: daemon.Activate})))
+}

pkg/buildinfo/buildinfo.go

@@ -26,9 +26,6 @@ var VersionSuffix = "-dev.unknown"
 // overridden when building Elvish; see PACKAGING.md for details.
 var Reproducible = "false"
 
-// Program is the buildinfo subprogram.
-var Program prog.Program = program{}
-
 // Type contains all the build information fields.
 type Type struct {
 	Version      string `json:"version"`
@@ -45,26 +42,36 @@ var Value = Type{
 	GoVersion:    runtime.Version(),
 }
 
-type program struct{}
+// Program is the buildinfo subprogram.
+type Program struct {
+	version, buildinfo bool
+	json               *bool
+}
+
+func (p *Program) RegisterFlags(fs *prog.FlagSet) {
+	fs.BoolVar(&p.version, "version", false, "show version and quit")
+	fs.BoolVar(&p.buildinfo, "buildinfo", false, "show build info and quit")
+	p.json = fs.JSON()
+}
 
-func (program) Run(fds [3]*os.File, f *prog.Flags, _ []string) error {
+func (p *Program) Run(fds [3]*os.File, _ []string) error {
 	switch {
-	case f.BuildInfo:
-		if f.JSON {
+	case p.buildinfo:
+		if *p.json {
 			fmt.Fprintln(fds[1], mustToJSON(Value))
 		} else {
 			fmt.Fprintln(fds[1], "Version:", Value.Version)
 			fmt.Fprintln(fds[1], "Go version:", Value.GoVersion)
 			fmt.Fprintln(fds[1], "Reproducible build:", Value.Reproducible)
 		}
-	case f.Version:
-		if f.JSON {
+	case p.version:
+		if *p.json {
 			fmt.Fprintln(fds[1], mustToJSON(Value.Version))
 		} else {
 			fmt.Fprintln(fds[1], Value.Version)
 		}
 	default:
-		return prog.ErrNotSuitable
+		return prog.ErrNextProgram
 	}
 	return nil
 }

pkg/buildinfo/buildinfo_test.go

@@ -8,7 +8,7 @@ import (
 )
 
 func TestProgram(t *testing.T) {
-	Test(t, Program,
+	Test(t, &Program{},
 		ThatElvish("-version").WritesStdout(Value.Version+"\n"),
 		ThatElvish("-version", "-json").WritesStdout(mustToJSON(Value.Version)+"\n"),
 

pkg/daemon/server.go

@@ -21,21 +21,31 @@ import (
 var logger = logutil.GetLogger("[daemon] ")
 
 // Program is the daemon subprogram.
-var Program prog.Program = program{}
+type Program struct {
+	run   bool
+	paths *prog.DaemonPaths
+	// Used in tests.
+	serveOpts ServeOpts
+}
 
-type program struct {
-	ServeOpts ServeOpts
+func (p *Program) RegisterFlags(fs *prog.FlagSet) {
+	fs.BoolVar(&p.run, "daemon", false, "[internal flag] run the storage daemon instead of shell")
+	p.paths = fs.DaemonPaths()
 }
 
-func (p program) Run(fds [3]*os.File, f *prog.Flags, args []string) error {
-	if !f.Daemon {
-		return prog.ErrNotSuitable
+func (p *Program) Run(fds [3]*os.File, args []string) error {
+	if !p.run {
+		return prog.ErrNextProgram
 	}
 	if len(args) > 0 {
 		return prog.BadUsage("arguments are not allowed with -daemon")
 	}
+
+	// The stdout is redirected to a unique log file (see the spawn function),
+	// so just use it for logging.
+	logutil.SetOutput(fds[1])
 	setUmaskForDaemon()
-	exit := Serve(f.Sock, f.DB, p.ServeOpts)
+	exit := Serve(p.paths.Sock, p.paths.DB, p.serveOpts)
 	return prog.Exit(exit)
 }
 

pkg/daemon/server_test.go

@@ -17,7 +17,7 @@ func TestProgram_TerminatesIfCannotListen(t *testing.T) {
 	setup(t)
 	testutil.MustCreateEmpty("sock")
 
-	Test(t, Program,
+	Test(t, &Program{},
 		ThatElvish("-daemon", "-sock", "sock", "-db", "db").
 			ExitsWith(2).
 			WritesStdoutContaining("failed to listen on sock"),
@@ -82,7 +82,7 @@ func TestProgram_QuitsOnSignalChannelWithClients(t *testing.T) {
 }
 
 func TestProgram_BadCLI(t *testing.T) {
-	Test(t, Program,
+	Test(t, &Program{},
 		ThatElvish().
 			ExitsWith(2).
 			WritesStderr("internal error: no suitable subprogram\n"),
@@ -118,7 +118,7 @@ func startServerOpts(t *testing.T, args []string, opts ServeOpts) server {
 	opts.Ready = readyCh
 	doneCh := make(chan serverResult)
 	go func() {
-		exit, stdout, stderr := Run(program{opts}, args...)
+		exit, stdout, stderr := Run(&Program{serveOpts: opts}, args...)
 		doneCh <- serverResult{exit, stdout, stderr}
 		close(doneCh)
 	}()

pkg/lsp/lsp.go

@@ -10,13 +10,17 @@ import (
 )
 
 // Program is the LSP subprogram.
-var Program prog.Program = program{}
+type Program struct {
+	run bool
+}
 
-type program struct{}
+func (p *Program) RegisterFlags(fs *prog.FlagSet) {
+	fs.BoolVar(&p.run, "lsp", false, "run language server instead of shell")
+}
 
-func (program) Run(fds [3]*os.File, f *prog.Flags, _ []string) error {
-	if !f.LSP {
-		return prog.ErrNotSuitable
+func (p *Program) Run(fds [3]*os.File, _ []string) error {
+	if !p.run {
+		return prog.ErrNextProgram
 	}
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

pkg/lsp/lsp_test.go

@@ -113,7 +113,7 @@ func TestJSONRPCErrors(t *testing.T) {
 }
 
 func TestProgramErrors(t *testing.T) {
-	progtest.Test(t, Program,
+	progtest.Test(t, &Program{},
 		progtest.ThatElvish("").
 			ExitsWith(2).
 			WritesStderr("internal error: no suitable subprogram\n"))
@@ -164,7 +164,7 @@ func setup(t *testing.T) *clientFixture {
 	// Run server
 	done := make(chan struct{})
 	go func() {
-		prog.Run([3]*os.File{r0, w1, nil}, []string{"elvish", "-lsp"}, Program)
+		prog.Run([3]*os.File{r0, w1, nil}, []string{"elvish", "-lsp"}, &Program{})
 		close(done)
 	}()
 	t.Cleanup(func() { <-done })

pkg/pprof/pprof.go

@@ -0,0 +1,33 @@
+// Package pprof adds profiling support to the Elvish program.
+package pprof
+
+import (
+	"fmt"
+	"os"
+	"runtime/pprof"
+
+	"src.elv.sh/pkg/prog"
+)
+
+// Program adds support for the -cpuprofile flag.
+type Program struct {
+	cpuProfile string
+}
+
+func (p *Program) RegisterFlags(f *prog.FlagSet) {
+	f.StringVar(&p.cpuProfile, "cpuprofile", "", "write CPU profile to file")
+}
+
+func (p *Program) Run(fds [3]*os.File, _ []string) error {
+	if p.cpuProfile != "" {
+		f, err := os.Create(p.cpuProfile)
+		if err != nil {
+			fmt.Fprintln(fds[2], "Warning: cannot create CPU profile:", err)
+			fmt.Fprintln(fds[2], "Continuing without CPU profiling.")
+		} else {
+			pprof.StartCPUProfile(f)
+			defer pprof.StopCPUProfile()
+		}
+	}
+	return prog.ErrNextProgram
+}

pkg/pprof/pprof_test.go

@@ -0,0 +1,38 @@
+package pprof_test
+
+import (
+	"os"
+	"testing"
+
+	. "src.elv.sh/pkg/pprof"
+	"src.elv.sh/pkg/prog"
+	"src.elv.sh/pkg/prog/progtest"
+	"src.elv.sh/pkg/testutil"
+)
+
+var (
+	Test       = progtest.Test
+	ThatElvish = progtest.ThatElvish
+)
+
+func TestProgram(t *testing.T) {
+	testutil.InTempDir(t)
+
+	Test(t, prog.Composite(&Program{}, noopProgram{}),
+		ThatElvish("-cpuprofile", "cpuprof").DoesNothing(),
+		ThatElvish("-cpuprofile", "/a/bad/path").
+			WritesStderrContaining("Warning: cannot create CPU profile:"),
+	)
+
+	// Check for the effect of -cpuprofile. There isn't much to test beyond a
+	// sanity check that the profile file now exists.
+	_, err := os.Stat("cpuprof")
+	if err != nil {
+		t.Errorf("CPU profile file does not exist: %v", err)
+	}
+}
+
+type noopProgram struct{}
+
+func (noopProgram) RegisterFlags(*prog.FlagSet)     {}
+func (noopProgram) Run([3]*os.File, []string) error { return nil }

pkg/prog/flags.go

@@ -0,0 +1,32 @@
+package prog
+
+import "flag"
+
+type FlagSet struct {
+	*flag.FlagSet
+	daemonPaths *DaemonPaths
+	json        *bool
+}
+
+type DaemonPaths struct {
+	DB, Sock string
+}
+
+func (fs *FlagSet) DaemonPaths() *DaemonPaths {
+	if fs.daemonPaths == nil {
+		var dp DaemonPaths
+		fs.StringVar(&dp.DB, "db", "", "[internal flag] path to the database")
+		fs.StringVar(&dp.Sock, "sock", "", "[internal flag] path to the daemon socket")
+		fs.daemonPaths = &dp
+	}
+	return fs.daemonPaths
+}
+
+func (fs *FlagSet) JSON() *bool {
+	if fs.json == nil {
+		var json bool
+		fs.BoolVar(&json, "json", false, "show output in JSON. Useful with -buildinfo and -compileonly")
+		fs.json = &json
+	}
+	return fs.json
+}

pkg/prog/prog.go

@@ -12,7 +12,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"runtime/pprof"
 
 	"src.elv.sh/pkg/logutil"
 )
@@ -22,52 +21,11 @@ import (
 // 0.X.
 var DeprecationLevel = 17
 
-// Flags keeps command-line flags.
-type Flags struct {
-	Log, CPUProfile string
-
-	Help, Version, BuildInfo, JSON bool
-
-	CodeInArg, CompileOnly, NoRc bool
-	RC                           string
-
-	Daemon   bool
-	DB, Sock string
-
-	LSP bool
-}
-
-func newFlagSet(f *Flags) *flag.FlagSet {
-	fs := flag.NewFlagSet("elvish", flag.ContinueOnError)
-	// Error and usage will be printed explicitly.
-	fs.SetOutput(io.Discard)
-
-	fs.StringVar(&f.Log, "log", "", "a file to write debug log to except for the daemon")
-	fs.StringVar(&f.CPUProfile, "cpuprofile", "", "write cpu profile to file")
-
-	fs.BoolVar(&f.Help, "help", false, "show usage help and quit")
-	fs.BoolVar(&f.Version, "version", false, "show version and quit")
-	fs.BoolVar(&f.BuildInfo, "buildinfo", false, "show build info and quit")
-	fs.BoolVar(&f.JSON, "json", false, "show output in JSON. Useful with -buildinfo and -compileonly")
-
-	// The `-i` option is for compatibility with POSIX shells so that programs, such as the `script`
-	// command, will work when asked to launch an interactive Elvish shell.
-	fs.Bool("i", false, "force interactive mode; currently ignored")
-	fs.BoolVar(&f.CodeInArg, "c", false, "take first argument as code to execute")
-	fs.BoolVar(&f.CompileOnly, "compileonly", false, "Parse/Compile but do not execute")
-	fs.BoolVar(&f.NoRc, "norc", false, "run elvish without invoking rc.elv")
-	fs.StringVar(&f.RC, "rc", "", "path to rc.elv")
-
-	fs.BoolVar(&f.Daemon, "daemon", false, "[internal flag] run the storage daemon instead of shell")
-
-	fs.StringVar(&f.DB, "db", "", "[internal flag] path to the database")
-	fs.StringVar(&f.Sock, "sock", "", "[internal flag] path to the daemon socket")
-
-	fs.BoolVar(&f.LSP, "lsp", false, "run LSP server instead of shell")
-
-	fs.IntVar(&DeprecationLevel, "deprecation-level", DeprecationLevel, "show warnings for all features deprecated as of version 0.X")
-
-	return fs
+// Program represents a subprogram.
+type Program interface {
+	RegisterFlags(fs *FlagSet)
+	// Run runs the subprogram.
+	Run(fds [3]*os.File, args []string) error
 }
 
 func usage(out io.Writer, fs *flag.FlagSet) {
@@ -80,8 +38,18 @@ func usage(out io.Writer, fs *flag.FlagSet) {
 // Run parses command-line flags and runs the first applicable subprogram. It
 // returns the exit status of the program.
 func Run(fds [3]*os.File, args []string, p Program) int {
-	f := &Flags{}
-	fs := newFlagSet(f)
+	fs := flag.NewFlagSet("elvish", flag.ContinueOnError)
+	// Error and usage will be printed explicitly.
+	fs.SetOutput(io.Discard)
+
+	var log string
+	var help bool
+	fs.StringVar(&log, "log", "", "a file to write debug log to except for the daemon")
+	fs.BoolVar(&help, "help", false, "show usage help and quit")
+	fs.IntVar(&DeprecationLevel, "deprecation-level", DeprecationLevel, "show warnings for all features deprecated as of version 0.X")
+
+	p.RegisterFlags(&FlagSet{FlagSet: fs})
+
 	err := fs.Parse(args[1:])
 	if err != nil {
 		if err == flag.ErrHelp {
@@ -97,38 +65,25 @@ func Run(fds [3]*os.File, args []string, p Program) int {
 		return 2
 	}
 
-	// Handle flags common to all subprograms.
-	if f.CPUProfile != "" {
-		f, err := os.Create(f.CPUProfile)
-		if err != nil {
-			fmt.Fprintln(fds[2], "Warning: cannot create CPU profile:", err)
-			fmt.Fprintln(fds[2], "Continuing without CPU profiling.")
-		} else {
-			pprof.StartCPUProfile(f)
-			defer pprof.StopCPUProfile()
-		}
-	}
-
-	if f.Daemon {
-		// We expect our stdout file handle is open on a unique log file for the daemon to write its
-		// log messages. See daemon.Spawn() in pkg/daemon.
-		logutil.SetOutput(fds[1])
-	} else if f.Log != "" {
-		err = logutil.SetOutputFile(f.Log)
+	if log != "" {
+		err = logutil.SetOutputFile(log)
 		if err != nil {
 			fmt.Fprintln(fds[2], err)
 		}
 	}
 
-	if f.Help {
+	if help {
 		usage(fds[1], fs)
 		return 0
 	}
 
-	err = p.Run(fds, f, fs.Args())
+	err = p.Run(fds, fs.Args())
 	if err == nil {
 		return 0
 	}
+	if err == ErrNextProgram {
+		err = errNoSuitableSubprogram
+	}
 	if msg := err.Error(); msg != "" {
 		fmt.Fprintln(fds[2], msg)
 	}
@@ -144,26 +99,34 @@ func Run(fds [3]*os.File, args []string, p Program) int {
 // Composite returns a Program that tries each of the given programs,
 // terminating at the first one that doesn't return NotSuitable().
 func Composite(programs ...Program) Program {
-	return compositeProgram(programs)
+	return composite(programs)
 }
 
-type compositeProgram []Program
+type composite []Program
+
+func (cp composite) RegisterFlags(f *FlagSet) {
+	for _, p := range cp {
+		p.RegisterFlags(f)
+	}
+}
 
-func (cp compositeProgram) Run(fds [3]*os.File, f *Flags, args []string) error {
+func (cp composite) Run(fds [3]*os.File, args []string) error {
 	for _, p := range cp {
-		err := p.Run(fds, f, args)
-		if err != ErrNotSuitable {
+		err := p.Run(fds, args)
+		if err != ErrNextProgram {
 			return err
 		}
 	}
-	// If we have reached here, all subprograms have returned errNotSuitable
-	return ErrNotSuitable
+	// If we have reached here, all subprograms have returned ErrNextProgram
+	return ErrNextProgram
 }
 
-// ErrNotSuitable is a special error that may be returned by Program.Run, to
-// signify that this Program should not be run. It is useful when a Program is
-// used in Composite.
-var ErrNotSuitable = errors.New("internal error: no suitable subprogram")
+var errNoSuitableSubprogram = errors.New("internal error: no suitable subprogram")
+
+// ErrNextProgram is a special error that may be returned by Program.Run that
+// is part of a Composite program, indicating that the next program should be
+// tried.
+var ErrNextProgram = errors.New("next program")
 
 // BadUsage returns a special error that may be returned by Program.Run. It
 // causes the main function to print out a message, the usage information and
@@ -187,9 +150,3 @@ func Exit(exit int) error {
 type exitError struct{ exit int }
 
 func (e exitError) Error() string { return "" }
-
-// Program represents a subprogram.
-type Program interface {
-	// Run runs the subprogram.
-	Run(fds [3]*os.File, f *Flags, args []string) error
-}

pkg/prog/prog_test.go

@@ -1,6 +1,7 @@
 package prog_test
 
 import (
+	"fmt"
 	"os"
 	"testing"
 
@@ -14,10 +15,8 @@ var (
 	ThatElvish = progtest.ThatElvish
 )
 
-func TestCommonFlagHandling(t *testing.T) {
-	testutil.InTempDir(t)
-
-	Test(t, testProgram{},
+func TestFlagHandling(t *testing.T) {
+	Test(t, &testProgram{},
 		ThatElvish("-bad-flag").
 			ExitsWith(2).
 			WritesStderrContaining("flag provided but not defined: -bad-flag\nUsage:"),
@@ -28,24 +27,47 @@ func TestCommonFlagHandling(t *testing.T) {
 
 		ThatElvish("-help").
 			WritesStdoutContaining("Usage: elvish [flags] [script]"),
-
-		ThatElvish("-cpuprofile", "cpuprof").DoesNothing(),
-		ThatElvish("-cpuprofile", "/a/bad/path").
-			WritesStderrContaining("Warning: cannot create CPU profile:"),
 	)
+}
 
-	// Check for the effect of -cpuprofile. There isn't much to test beyond a
-	// sanity check that the profile file now exists.
-	_, err := os.Stat("cpuprof")
+func TestLogFlag(t *testing.T) {
+	testutil.InTempDir(t)
+	Test(t, &testProgram{},
+		ThatElvish("-log", "log").DoesNothing())
+	_, err := os.Stat("log")
 	if err != nil {
-		t.Errorf("CPU profile file does not exist: %v", err)
+		t.Errorf("log file was not created: %v", err)
 	}
 }
 
+func TestCustomFlag(t *testing.T) {
+	Test(t, &testProgram{customFlag: true},
+		ThatElvish("-flag", "foo").
+			WritesStdout("-flag foo\n"),
+	)
+}
+
+func TestSharedFlags(t *testing.T) {
+	Test(t, &testProgram{sharedFlags: true},
+		ThatElvish("-sock", "sock", "-db", "db", "-json").
+			WritesStdout("-sock sock -db db -json true\n"),
+	)
+}
+
+func TestSharedFlags_MultiplePrograms(t *testing.T) {
+	Test(t,
+		Composite(
+			&testProgram{sharedFlags: true, nextProgram: true},
+			&testProgram{sharedFlags: true}),
+		ThatElvish("-sock", "sock", "-db", "db", "-json").
+			WritesStdout("-sock sock -db db -json true\n"),
+	)
+}
+
 func TestShowDeprecations(t *testing.T) {
 	progtest.SetDeprecationLevel(t, 0)
 
-	Test(t, testProgram{},
+	Test(t, &testProgram{},
 		ThatElvish("-deprecation-level", "42").DoesNothing(),
 	)
 
@@ -55,7 +77,7 @@ func TestShowDeprecations(t *testing.T) {
 }
 
 func TestNoSuitableSubprogram(t *testing.T) {
-	Test(t, testProgram{notSuitable: true},
+	Test(t, &testProgram{nextProgram: true},
 		ThatElvish().
 			ExitsWith(2).
 			WritesStderr("internal error: no suitable subprogram\n"),
@@ -64,14 +86,14 @@ func TestNoSuitableSubprogram(t *testing.T) {
 
 func TestComposite(t *testing.T) {
 	Test(t,
-		Composite(testProgram{notSuitable: true}, testProgram{writeOut: "program 2"}),
+		Composite(&testProgram{nextProgram: true}, &testProgram{writeOut: "program 2"}),
 		ThatElvish().WritesStdout("program 2"),
 	)
 }
 
 func TestComposite_NoSuitableSubprogram(t *testing.T) {
 	Test(t,
-		Composite(testProgram{notSuitable: true}, testProgram{notSuitable: true}),
+		Composite(&testProgram{nextProgram: true}, &testProgram{nextProgram: true}),
 		ThatElvish().
 			ExitsWith(2).
 			WritesStderr("internal error: no suitable subprogram\n"),
@@ -81,40 +103,63 @@ func TestComposite_NoSuitableSubprogram(t *testing.T) {
 func TestComposite_PreferEarlierSubprogram(t *testing.T) {
 	Test(t,
 		Composite(
-			testProgram{writeOut: "program 1"}, testProgram{writeOut: "program 2"}),
+			&testProgram{writeOut: "program 1"}, &testProgram{writeOut: "program 2"}),
 		ThatElvish().WritesStdout("program 1"),
 	)
 }
 
 func TestBadUsageError(t *testing.T) {
 	Test(t,
-		testProgram{returnErr: BadUsage("lorem ipsum")},
+		&testProgram{returnErr: BadUsage("lorem ipsum")},
 		ThatElvish().ExitsWith(2).WritesStderrContaining("lorem ipsum\n"),
 	)
 }
 
 func TestExitError(t *testing.T) {
-	Test(t, testProgram{returnErr: Exit(3)},
+	Test(t, &testProgram{returnErr: Exit(3)},
 		ThatElvish().ExitsWith(3),
 	)
 }
 
 func TestExitError_0(t *testing.T) {
-	Test(t, testProgram{returnErr: Exit(0)},
+	Test(t, &testProgram{returnErr: Exit(0)},
 		ThatElvish().ExitsWith(0),
 	)
 }
 
 type testProgram struct {
-	notSuitable bool
+	nextProgram bool
 	writeOut    string
 	returnErr   error
+	customFlag  bool
+	sharedFlags bool
+
+	flag        string
+	daemonPaths *DaemonPaths
+	json        *bool
 }
 
-func (p testProgram) Run(fds [3]*os.File, _ *Flags, args []string) error {
-	if p.notSuitable {
-		return ErrNotSuitable
+func (p *testProgram) RegisterFlags(f *FlagSet) {
+	if p.customFlag {
+		f.StringVar(&p.flag, "flag", "default", "a flag")
+	}
+	if p.sharedFlags {
+		p.daemonPaths = f.DaemonPaths()
+		p.json = f.JSON()
+	}
+}
+
+func (p *testProgram) Run(fds [3]*os.File, args []string) error {
+	if p.nextProgram {
+		return ErrNextProgram
 	}
 	fds[1].WriteString(p.writeOut)
+	if p.customFlag {
+		fmt.Fprintf(fds[1], "-flag %s\n", p.flag)
+	}
+	if p.sharedFlags {
+		fmt.Fprintf(fds[1], "-sock %s -db %s -json %v\n",
+			p.daemonPaths.Sock, p.daemonPaths.DB, *p.json)
+	}
 	return p.returnErr
 }

pkg/prog/progtest/progtest_test.go

@@ -17,7 +17,9 @@ func TestOutputCaptureDoesNotDeadlock(t *testing.T) {
 
 type noisyProgram struct{}
 
-func (noisyProgram) Run(fds [3]*os.File, f *prog.Flags, args []string) error {
+func (noisyProgram) RegisterFlags(f *prog.FlagSet) {}
+
+func (noisyProgram) Run(fds [3]*os.File, args []string) error {
 	// We need enough data to verify whether we're likely to deadlock due to
 	// filling the pipe before the test completes. Pipes typically buffer 8 to
 	// 128 KiB.

pkg/shell/interact_test.go

@@ -15,7 +15,7 @@ func TestInteract(t *testing.T) {
 	MustWriteFile("rc-dnc.elv", "echo $a")
 	MustWriteFile("rc-fail.elv", "fail bad")
 
-	Test(t, Program{},
+	Test(t, &Program{},
 		thatElvishInteract().WithStdin("echo hello\n").WritesStdout("hello\n"),
 		thatElvishInteract().WithStdin("fail mock\n").WritesStderrContaining("fail mock"),
 
@@ -37,7 +37,7 @@ func TestInteract_DefaultRCPath(t *testing.T) {
 		filepath.Join(home, ".elvish", "rc.elv"), "echo hello legacy rc.elv")
 	// Note: non-legacy path is tested in interact_unix_test.go
 
-	Test(t, Program{},
+	Test(t, &Program{},
 		thatElvishInteract().WritesStdout("hello legacy rc.elv\n"),
 	)
 }

pkg/shell/interact_unix_test.go

@@ -22,7 +22,7 @@ func TestInteract_NewRcFile_Default(t *testing.T) {
 	MustWriteFile(
 		filepath.Join(home, ".config", "elvish", "rc.elv"), "echo hello new rc.elv")
 
-	Test(t, Program{},
+	Test(t, &Program{},
 		thatElvishInteract().WritesStdout("hello new rc.elv\n"),
 	)
 }
@@ -34,7 +34,7 @@ func TestInteract_NewRcFile_XDG_CONFIG_HOME(t *testing.T) {
 		filepath.Join(xdgConfigHome, "elvish", "rc.elv"),
 		"echo hello XDG_CONFIG_HOME rc.elv")
 
-	Test(t, Program{},
+	Test(t, &Program{},
 		thatElvishInteract().WritesStdout("hello XDG_CONFIG_HOME rc.elv\n"),
 	)
 }
@@ -63,7 +63,7 @@ func TestInteract_ConnectsToDaemon(t *testing.T) {
 		t.Fatalf("timed out waiting for daemon to start")
 	}
 
-	Test(t, Program{daemon.Activate},
+	Test(t, &Program{ActivateDaemon: daemon.Activate},
 		thatElvishInteract("-sock", "sock", "-db", "db").
 			WithStdin("use daemon; echo $daemon:pid\n").
 			WritesStdout(fmt.Sprintln(os.Getpid())),

pkg/shell/paths.go

@@ -26,17 +26,17 @@ func libPaths() ([]string, error) {
 
 // Returns a SpawnConfig containing all the paths needed by the daemon. It
 // respects overrides of sock and db from CLI flags.
-func daemonPaths(flags *prog.Flags) (*daemondefs.SpawnConfig, error) {
+func daemonPaths(p *prog.DaemonPaths) (*daemondefs.SpawnConfig, error) {
 	runDir, err := secureRunDir()
 	if err != nil {
 		return nil, err
 	}
-	sock := flags.Sock
+	sock := p.Sock
 	if sock == "" {
 		sock = filepath.Join(runDir, "sock")
 	}
 
-	db := flags.DB
+	db := p.DB
 	if db == "" {
 		var err error
 		db, err = dbPath()

pkg/shell/script_test.go

@@ -11,7 +11,7 @@ func TestScript(t *testing.T) {
 	testutil.InTempDir(t)
 	testutil.MustWriteFile("a.elv", "echo hello")
 
-	Test(t, Program{},
+	Test(t, &Program{},
 		ThatElvish("a.elv").WritesStdout("hello\n"),
 		ThatElvish("-c", "echo hello").WritesStdout("hello\n"),
 		ThatElvish("non-existent.elv").

pkg/shell/shell.go

@@ -25,9 +25,31 @@ var logger = logutil.GetLogger("[shell] ")
 // Program is the shell subprogram.
 type Program struct {
 	ActivateDaemon daemondefs.ActivateFunc
+
+	codeInArg   bool
+	compileOnly bool
+	noRC        bool
+	rc          string
+	json        *bool
+	daemonPaths *prog.DaemonPaths
+}
+
+func (p *Program) RegisterFlags(fs *prog.FlagSet) {
+	// Support -i so that programs that expect shells to support it (like
+	// "script") don't error when they invoke Elvish.
+	fs.Bool("i", false, "force interactive mode; currently ignored")
+	fs.BoolVar(&p.codeInArg, "c", false, "take first argument as code to execute")
+	fs.BoolVar(&p.compileOnly, "compileonly", false, "Parse/Compile but do not execute")
+	fs.BoolVar(&p.noRC, "norc", false, "run elvish without invoking rc.elv")
+	fs.StringVar(&p.rc, "rc", "", "path to rc.elv")
+
+	p.json = fs.JSON()
+	if p.ActivateDaemon != nil {
+		p.daemonPaths = fs.DaemonPaths()
+	}
 }
 
-func (p Program) Run(fds [3]*os.File, f *prog.Flags, args []string) error {
+func (p *Program) Run(fds [3]*os.File, args []string) error {
 	cleanup1 := IncSHLVL()
 	defer cleanup1()
 	cleanup2 := initTTYAndSignal(fds[2])
@@ -38,14 +60,14 @@ func (p Program) Run(fds [3]*os.File, f *prog.Flags, args []string) error {
 	if len(args) > 0 {
 		exit := script(
 			ev, fds, args, &scriptCfg{
-				Cmd: f.CodeInArg, CompileOnly: f.CompileOnly, JSON: f.JSON})
+				Cmd: p.codeInArg, CompileOnly: p.compileOnly, JSON: *p.json})
 		return prog.Exit(exit)
 	}
 
 	var spawnCfg *daemondefs.SpawnConfig
 	if p.ActivateDaemon != nil {
 		var err error
-		spawnCfg, err = daemonPaths(f)
+		spawnCfg, err = daemonPaths(p.daemonPaths)
 		if err != nil {
 			fmt.Fprintln(fds[2], "Warning:", err)
 			fmt.Fprintln(fds[2], "Storage daemon may not function.")
@@ -54,11 +76,11 @@ func (p Program) Run(fds [3]*os.File, f *prog.Flags, args []string) error {
 
 	rc := ""
 	switch {
-	case f.NoRc:
-	// Leave rc empty
-	case f.RC != "":
+	case p.noRC:
+		// Leave rc empty
+	case p.rc != "":
 		// Use explicit -rc flag value
-		rc = f.RC
+		rc = p.rc
 	default:
 		// Use default path to rc.elv
 		var err error

pkg/shell/shell_test.go

@@ -14,7 +14,7 @@ func TestShell_LegacyLibPath(t *testing.T) {
 	home := setupHomePaths(t)
 	MustWriteFile(filepath.Join(home, ".elvish", "lib", "a.elv"), "echo mod a")
 
-	Test(t, Program{},
+	Test(t, &Program{},
 		ThatElvish("-c", "use a").WritesStdout("mod a\n"),
 	)
 }