1// ================================================================
2// This handles print, printn, eprint, and eprintn statements.
3// ================================================================
4
5package cst
6
7import (
8	"bytes"
9	"errors"
10	"fmt"
11	"os"
12
13	"miller/src/dsl"
14	"miller/src/lib"
15	"miller/src/output"
16	"miller/src/runtime"
17	"miller/src/types"
18)
19
20// ----------------------------------------------------------------
21// Example ASTs:
22//
23// $ mlr -n put -v 'print $a, $b'
24// DSL EXPRESSION:
25// print $a, $b
26// AST:
27// * statement block
28//     * print statement "print"
29//         * function callsite
30//             * direct field value "a"
31//             * direct field value "b"
32//         * no-op
33//
34// $ mlr -n put -v 'print > stdout, $a, $b'
35// DSL EXPRESSION:
36// print > stdout, $a, $b
37// AST:
38// * statement block
39//     * print statement "print"
40//         * function callsite
41//             * direct field value "a"
42//             * direct field value "b"
43//         * redirect write ">"
44//             * stdout redirect target "stdout"
45//
46// $ mlr -n put -v 'print > stderr, $a, $b'
47// DSL EXPRESSION:
48// print > stderr, $a, $b
49// AST:
50// * statement block
51//     * print statement "print"
52//         * function callsite
53//             * direct field value "a"
54//             * direct field value "b"
55//         * redirect write ">"
56//             * stderr redirect target "stderr"
57//
58// $ mlr -n put -v 'print > "foo.dat", $a, $b'
59// DSL EXPRESSION:
60// print > "foo.dat", $a, $b
61// AST:
62// * statement block
63//     * print statement "print"
64//         * function callsite
65//             * direct field value "a"
66//             * direct field value "b"
67//         * redirect write ">"
68//             * string literal "foo.dat"
69//
70// $ mlr -n put -v 'print >> "foo.dat", $a, $b'
71// DSL EXPRESSION:
72// print >> "foo.dat", $a, $b
73// AST:
74// * statement block
75//     * print statement "print"
76//         * function callsite
77//             * direct field value "a"
78//             * direct field value "b"
79//         * redirect append ">>"
80//             * string literal "foo.dat"
81//
82// $ mlr -n put -v 'print | "command", $a, $b'
83// DSL EXPRESSION:
84// print | "command", $a, $b
85// AST:
86// * statement block
87//     * print statement "print"
88//         * function callsite
89//             * direct field value "a"
90//             * direct field value "b"
91//         * redirect pipe "|"
92//             * string literal "command"
93//
94//  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
95// Corresponding data structures for these cases:
96//
97// * printToRedirectFunc is either printToStdout, printToStderr, or
98//   printToFileOrPipe. Only the third of these takes a non-nil
99//   redirectorTargetEvaluable and a non-nil outputHandlerManager.
100//
101// * redirectorTargetEvaluable is nil for stdout or stderr.
102//
103// * The OutputHandlerManager is for file names or commands in >, >> or |.
104//   This is because the target for the redirect can vary from one record to
105//   the next, e.g. mlr put 'print > $a.txt, $b'. The OutputHandlerManager
106//   keeps file-handles for each distinct value of $a.
107//
108// So:
109//
110// * print $a, $b
111//   AST redirectorNode         = NodeTypeNoOp
112//     AST redirectorTargetNode = (none)
113//   printToRedirectFunc        = printToStdout
114//   redirectorTargetEvaluable  = nil
115//   outputHandlerManager       = nil
116//
117// * print > stdout, $a, $b
118//   AST redirectorNode         = NodeTypeRedirectWrite
119//     AST redirectorTargetNode = NodeTypeRedirectTargetStdout
120//   printToRedirectFunc        = printToStdout
121//   redirectorTargetEvaluable  = nil
122//   outputHandlerManager       = nil
123//
124// * print > stderr, $a, $b
125//   AST redirectorNode         = NodeTypeRedirectWrite
126//     AST redirectorTargetNode = NodeTypeRedirectTargetStderr
127//   printToRedirectFunc        = printToStderr
128//   redirectorTargetEvaluable  = nil
129//   outputHandlerManager       = nil
130//
131// * print > "foo.dat", $a, $b
132//   AST redirectorNode         = NodeTypeRedirectWrite
133//     AST redirectorTargetNode = any of various evaluables
134//   printToRedirectFunc        = printToFileOrPipe
135//   redirectorTargetEvaluable  = non-nil
136//   outputHandlerManager       = non-nil
137//
138// * print >> "foo.dat", $a, $b
139//   AST redirectorNode         = NodeTypeRedirectAppend
140//     AST redirectorTargetNode = any of various evaluables
141//   printToRedirectFunc        = printToFileOrPipe
142//   redirectorTargetEvaluable  = non-nil
143//   outputHandlerManager       = non-nil
144//
145// * print | "command", $a, $b
146//   AST redirectorNode         = NodeTypeRedirectPipe
147//     AST redirectorTargetNode = any of various evaluables
148//   printToRedirectFunc        = printToFileOrPipe
149//   redirectorTargetEvaluable  = non-nil
150//   outputHandlerManager       = non-nil
151
152// ================================================================
153type tPrintToRedirectFunc func(
154	outputString string,
155	state *runtime.State,
156) error
157
158type PrintStatementNode struct {
159	expressionEvaluables      []IEvaluable
160	terminator                string
161	printToRedirectFunc       tPrintToRedirectFunc
162	redirectorTargetEvaluable IEvaluable                  // for file/pipe targets
163	outputHandlerManager      output.OutputHandlerManager // for file/pipe targets
164}
165
166// ----------------------------------------------------------------
167func (this *RootNode) BuildPrintStatementNode(astNode *dsl.ASTNode) (IExecutable, error) {
168	lib.InternalCodingErrorIf(astNode.Type != dsl.NodeTypePrintStatement)
169	return this.buildPrintxStatementNode(
170		astNode,
171		os.Stdout,
172		"\n",
173	)
174}
175
176func (this *RootNode) BuildPrintnStatementNode(astNode *dsl.ASTNode) (IExecutable, error) {
177	lib.InternalCodingErrorIf(astNode.Type != dsl.NodeTypePrintnStatement)
178	return this.buildPrintxStatementNode(
179		astNode,
180		os.Stdout,
181		"",
182	)
183}
184
185func (this *RootNode) BuildEprintStatementNode(astNode *dsl.ASTNode) (IExecutable, error) {
186	lib.InternalCodingErrorIf(astNode.Type != dsl.NodeTypeEprintStatement)
187	return this.buildPrintxStatementNode(
188		astNode,
189		os.Stderr,
190		"\n",
191	)
192}
193
194func (this *RootNode) BuildEprintnStatementNode(astNode *dsl.ASTNode) (IExecutable, error) {
195	lib.InternalCodingErrorIf(astNode.Type != dsl.NodeTypeEprintnStatement)
196	return this.buildPrintxStatementNode(
197		astNode,
198		os.Stderr,
199		"",
200	)
201}
202
203// ----------------------------------------------------------------
204// Common code for building print/eprint/printn/eprintn nodes
205
206func (this *RootNode) buildPrintxStatementNode(
207	astNode *dsl.ASTNode,
208	defaultOutputStream *os.File,
209	terminator string,
210) (IExecutable, error) {
211	lib.InternalCodingErrorIf(len(astNode.Children) != 2)
212	expressionsNode := astNode.Children[0]
213	redirectorNode := astNode.Children[1]
214
215	//  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
216	// Things to be printed, e.g. $a and $b in 'print > "foo.dat", $a, $b'.
217
218	var expressionEvaluables []IEvaluable = nil
219
220	if expressionsNode.Type == dsl.NodeTypeNoOp {
221		// Just 'print' without 'print $something'
222		expressionEvaluables = make([]IEvaluable, 1)
223		expressionEvaluable := this.BuildStringLiteralNode("")
224		expressionEvaluables[0] = expressionEvaluable
225	} else if expressionsNode.Type == dsl.NodeTypeFunctionCallsite {
226		expressionEvaluables = make([]IEvaluable, len(expressionsNode.Children))
227		for i, childNode := range expressionsNode.Children {
228			expressionEvaluable, err := this.BuildEvaluableNode(childNode)
229			if err != nil {
230				return nil, err
231			}
232			expressionEvaluables[i] = expressionEvaluable
233		}
234	} else {
235		lib.InternalCodingErrorIf(true)
236	}
237
238	//  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
239	// Redirection targets (the thing after > >> |, if any).
240
241	retval := &PrintStatementNode{
242		expressionEvaluables:      expressionEvaluables,
243		terminator:                terminator,
244		printToRedirectFunc:       nil,
245		redirectorTargetEvaluable: nil,
246		outputHandlerManager:      nil,
247	}
248
249	if redirectorNode.Type == dsl.NodeTypeNoOp {
250		// No > >> or | was provided.
251		if defaultOutputStream == os.Stdout {
252			retval.printToRedirectFunc = retval.printToStdout
253		} else if defaultOutputStream == os.Stderr {
254			retval.printToRedirectFunc = retval.printToStderr
255		} else {
256			lib.InternalCodingErrorIf(true)
257		}
258	} else {
259		// There is > >> or | provided.
260		lib.InternalCodingErrorIf(redirectorNode.Children == nil)
261		lib.InternalCodingErrorIf(len(redirectorNode.Children) != 1)
262		redirectorTargetNode := redirectorNode.Children[0]
263		var err error = nil
264
265		if redirectorTargetNode.Type == dsl.NodeTypeRedirectTargetStdout {
266			retval.printToRedirectFunc = retval.printToStdout
267		} else if redirectorTargetNode.Type == dsl.NodeTypeRedirectTargetStderr {
268			retval.printToRedirectFunc = retval.printToStderr
269		} else {
270			retval.printToRedirectFunc = retval.printToFileOrPipe
271
272			retval.redirectorTargetEvaluable, err = this.BuildEvaluableNode(redirectorTargetNode)
273			if err != nil {
274				return nil, err
275			}
276
277			if redirectorNode.Type == dsl.NodeTypeRedirectWrite {
278				retval.outputHandlerManager = output.NewFileWritetHandlerManager(this.recordWriterOptions)
279			} else if redirectorNode.Type == dsl.NodeTypeRedirectAppend {
280				retval.outputHandlerManager = output.NewFileAppendHandlerManager(this.recordWriterOptions)
281			} else if redirectorNode.Type == dsl.NodeTypeRedirectPipe {
282				retval.outputHandlerManager = output.NewPipeWriteHandlerManager(this.recordWriterOptions)
283			} else {
284				return nil, errors.New(
285					fmt.Sprintf(
286						"%s: unhandled redirector node type %s.",
287						lib.MlrExeName(), string(redirectorNode.Type),
288					),
289				)
290			}
291		}
292	}
293
294	// Register this with the CST root node so that open file descriptrs can be
295	// closed, etc at end of stream.
296	if retval.outputHandlerManager != nil {
297		this.RegisterOutputHandlerManager(retval.outputHandlerManager)
298	}
299
300	return retval, nil
301}
302
303// ----------------------------------------------------------------
304func (this *PrintStatementNode) Execute(state *runtime.State) (*BlockExitPayload, error) {
305	if len(this.expressionEvaluables) == 0 {
306		this.printToRedirectFunc(this.terminator, state)
307	} else {
308		// 5x faster than fmt.Print() separately: note that os.Stdout is
309		// non-buffered in Go whereas stdout is buffered in C.
310		//
311		// Minus: we need to do our own buffering for performance.
312		//
313		// Plus: we never have to worry about forgetting to do fflush(). :)
314		var buffer bytes.Buffer
315
316		for i, expressionEvaluable := range this.expressionEvaluables {
317			if i > 0 {
318				buffer.WriteString(" ")
319			}
320			evaluation := expressionEvaluable.Evaluate(state)
321			if !evaluation.IsAbsent() {
322				buffer.WriteString(evaluation.String())
323			}
324		}
325		buffer.WriteString(this.terminator)
326		this.printToRedirectFunc(buffer.String(), state)
327	}
328	return nil, nil
329}
330
331// ----------------------------------------------------------------
332func (this *PrintStatementNode) printToStdout(
333	outputString string,
334	state *runtime.State,
335) error {
336	// Insert the string into the record-output stream, so that goroutine can
337	// print it, resulting in deterministic output-ordering.
338
339	// The output channel is always non-nil, except for the Miller REPL.
340	if state.OutputChannel != nil {
341		state.OutputChannel <- types.NewOutputString(outputString, state.Context)
342	} else {
343		fmt.Print(outputString)
344	}
345
346	return nil
347}
348
349// ----------------------------------------------------------------
350func (this *PrintStatementNode) printToStderr(
351	outputString string,
352	state *runtime.State,
353) error {
354	fmt.Fprint(os.Stderr, outputString)
355	return nil
356}
357
358// ----------------------------------------------------------------
359func (this *PrintStatementNode) printToFileOrPipe(
360	outputString string,
361	state *runtime.State,
362) error {
363	redirectorTarget := this.redirectorTargetEvaluable.Evaluate(state)
364	if !redirectorTarget.IsString() {
365		return errors.New(
366			fmt.Sprintf(
367				"%s: output redirection yielded %s, not string.",
368				lib.MlrExeName(), redirectorTarget.GetTypeName(),
369			),
370		)
371	}
372	outputFileName := redirectorTarget.String()
373
374	this.outputHandlerManager.WriteString(outputString, outputFileName)
375	return nil
376}
377