37504: refactor code using/simulating memstream to capture output for "print -v"

1 files changed, 102 insertions, 71 deletions

diff --git a/Src/builtin.c b/Src/builtin.c
index cfc14a822..22011841d 100644
--- a/Src/builtin.c
+++ b/Src/builtin.c
@@ -4023,12 +4023,58 @@ bin_print(char *name, char **args, Options ops, int func)
     char *start, *endptr, *c, *d, *flag, *buf = NULL, spec[14], *fmt = NULL;
     char **first, **argp, *curarg, *flagch = "'0+- #", save = '\0', nullstr = '\0';
     size_t rcount, count = 0;
+    FILE *fout = stdout;
 #ifdef HAVE_OPEN_MEMSTREAM
     size_t mcount;
+#define ASSIGN_MSTREAM(BUF,FOUT) \
+    do { \
+        if ((fout = open_memstream(&BUF, &mcount)) == NULL) { \
+            zwarnnam(name, "open_memstream failed"); \
+            return 1; \
+        } \
+    } while (0)
+    /*
+     * Some implementations of open_memstream() have a bug such that,
+     * if fflush() is followed by fclose(), another NUL byte is written
+     * to the buffer at the wrong position.  Therefore we must fclose()
+     * before reading.
+     */
+#define READ_MSTREAM(BUF,COUNT,FOUT) \
+    (fclose(FOUT) == 0 ? (COUNT = mcount) : -1)
+#define CLOSE_MSTREAM(FOUT) 0
+
+#else /* simulate HAVE_OPEN_MEMSTREAM */
+
+#define ASSIGN_MSTREAM(BUF,FOUT) \
+    do { \
+        int tempfd; \
+        char *tmpf; \
+        if ((tempfd = gettempfile(NULL, 1, &tmpf)) < 0 || \
+            (fout = fdopen(tempfd, "w+")) == NULL) { \
+            zwarnnam(name, "can't open temp file: %e", errno); \
+            return 1; \
+        } \
+        unlink(tmpf); \
+    } while (0)
+#define READ_MSTREAM(BUF,COUNT,FOUT) \
+    ((((count = ftell(FOUT)), (BUF = (char *)zalloc(count + 1))) && \
+      ((fseek(FOUT, 0L, SEEK_SET) == 0) && !(BUF[count] = '\0')) && \
+      ((COUNT = fread(BUF, 1, count, FOUT)) == count)) ? count : -1)
+#define CLOSE_MSTREAM(FOUT) fclose(FOUT)
+
 #endif
-    FILE *fout = stdout;
-    Histent ent;
 
+#define IS_MSTREAM(FOUT) \
+    (FOUT != stdout && \
+     (OPT_ISSET(ops,'z') || OPT_ISSET(ops,'s') || OPT_ISSET(ops,'v')))
+
+    /* Testing EBADF special-cases >&- redirections */
+#define CLOSE_CLEANLY(FOUT) \
+    (IS_MSTREAM(FOUT) ? CLOSE_MSTREAM(FOUT) == 0 : \
+     ((FOUT == stdout) ? (fflush(FOUT) == 0 || errno == EBADF) : \
+      (fclose(FOUT) == 0)))	/* implies error for -u on a closed fd */
+
+    Histent ent;
     mnumber mnumval;
     double doubleval;
     int intval;
@@ -4227,6 +4273,10 @@ bin_print(char *name, char **args, Options ops, int func)
 	}
     }
 
+    if (OPT_ISSET(ops, 'v') ||
+	(fmt && (OPT_ISSET(ops,'z') || OPT_ISSET(ops,'s'))))
+	ASSIGN_MSTREAM(buf,fout);
+
     /* -c -- output in columns */
     if (!fmt && (OPT_ISSET(ops,'c') || OPT_ISSET(ops,'C'))) {
 	int l, nr, sc, n, t, i;
@@ -4378,12 +4428,22 @@ bin_print(char *name, char **args, Options ops, int func)
 	    }
 	    fputc(OPT_ISSET(ops,'N') ? '\0' : '\n', fout);
 	}
-	/* Testing EBADF special-cases >&- redirections */
-	if ((fout == stdout) ? (fflush(fout) != 0 && errno != EBADF) :
-	    (fclose(fout) != 0)) {
+	if (IS_MSTREAM(fout) && READ_MSTREAM(buf,rcount,fout) < 0)
+	    ret = 1;
+	if (!CLOSE_CLEANLY(fout) || ret) {
             zwarnnam(name, "write error: %e", errno);
             ret = 1;
 	}
+	if (buf) {
+	    /* assert: we must be doing -v at this point */
+	    queue_signals();
+	    if (ret)
+		free(buf);
+	    else
+		setsparam(OPT_ARG(ops, 'v'),
+			  metafy(buf, rcount, META_REALLOC));
+	    unqueue_signals();
+	}
 	return ret;
     }
 
@@ -4398,13 +4458,6 @@ bin_print(char *name, char **args, Options ops, int func)
 		metafy(args[n], len[n], META_NOALLOC);
 	}
 
-	/* -v option -- store the arguments in the named parameter */
-	if (OPT_ISSET(ops,'v')) {
-	    queue_signals();
-	    setsparam(OPT_ARG(ops, 'v'), sepjoin(args, NULL, 0));
-	    unqueue_signals();
-	    return 0;
-	}
 	/* -z option -- push the arguments onto the editing buffer stack */
 	if (OPT_ISSET(ops,'z')) {
 	    queue_signals();
@@ -4495,14 +4548,24 @@ bin_print(char *name, char **args, Options ops, int func)
 			  OPT_ISSET(ops,'N') ? '\0' : ' ', fout);
 	    }
 	}
-	if (!(OPT_ISSET(ops,'n') || nnl))
+	if (!(OPT_ISSET(ops,'n') || OPT_ISSET(ops, 'v') || nnl))
 	    fputc(OPT_ISSET(ops,'N') ? '\0' : '\n', fout);
-	/* Testing EBADF special-cases >&- redirections */
-	if ((fout == stdout) ? (fflush(fout) != 0 && errno != EBADF) :
-	    (fclose(fout) != 0)) {
+	if (IS_MSTREAM(fout) && READ_MSTREAM(buf,rcount,fout) < 0)
+	    ret = 1;
+	if (!CLOSE_CLEANLY(fout) || ret) {
             zwarnnam(name, "write error: %e", errno);
             ret = 1;
 	}
+	if (buf) {
+	    /* assert: we must be doing -v at this point */
+	    queue_signals();
+	    if (ret)
+		free(buf);
+	    else
+		setsparam(OPT_ARG(ops, 'v'),
+			  metafy(buf, rcount, META_REALLOC));
+	    unqueue_signals();
+	}
 	return ret;
     }
 
@@ -4512,20 +4575,6 @@ bin_print(char *name, char **args, Options ops, int func)
      * special cases of printing to a ZLE buffer or the history, however.
      */
 
-    if (OPT_ISSET(ops,'z') || OPT_ISSET(ops,'s') || OPT_ISSET(ops, 'v')) {
-#ifdef HAVE_OPEN_MEMSTREAM
-    	if ((fout = open_memstream(&buf, &mcount)) == NULL)
-	    zwarnnam(name, "open_memstream failed");
-#else
-	int tempfd;
-	char *tmpf;
-	if ((tempfd = gettempfile(NULL, 1, &tmpf)) < 0
-	 || (fout = fdopen(tempfd, "w+")) == NULL)
-	    zwarnnam(name, "can't open temp file: %e", errno);
-	unlink(tmpf);
-#endif
-    }
-
     /* printf style output */
     *spec = '%';
     argp = args;
@@ -4789,11 +4838,9 @@ bin_print(char *name, char **args, Options ops, int func)
 		}
 		zwarnnam(name, "%s: invalid directive", start);
 		if (*c) c[1] = save;
-		/* Testing EBADF special-cases >&- redirections */
-		if ((fout == stdout) ? (fflush(fout) != 0 && errno != EBADF) :
-		    (fclose(fout) != 0)) {
+		/* Why do we care about a clean close here? */
+		if (!CLOSE_CLEANLY(fout))
 		    zwarnnam(name, "write error: %e", errno);
-		}
 #ifdef HAVE_OPEN_MEMSTREAM
 		if (buf)
 		    free(buf);
@@ -4891,50 +4938,34 @@ bin_print(char *name, char **args, Options ops, int func)
 	/* if there are remaining args, reuse format string */
     } while (*argp && argp != first && !fmttrunc && !OPT_ISSET(ops,'r'));
 
-    if (OPT_ISSET(ops,'z') || OPT_ISSET(ops,'s') || OPT_ISSET(ops,'v')) {
-#ifdef HAVE_OPEN_MEMSTREAM
-	putc(0, fout);		/* not needed?  open_memstream() maintains? */
-	fclose(fout);
-	fout = NULL;
-	rcount = mcount;	/* now includes the trailing NUL we added */
-#else
-	rewind(fout);
-	buf = (char *)zalloc(count + 1);
-	rcount = fread(buf, 1, count, fout);
-	if (rcount < count)
-	    zwarnnam(name, "i/o error: %e", errno);
-	buf[rcount++] = '\0';
-#endif
+    if (IS_MSTREAM(fout)) {
 	queue_signals();
-	stringval = metafy(buf, rcount - 1, META_REALLOC);
-	buf = NULL;
-	if (OPT_ISSET(ops,'z')) {
-	    zpushnode(bufstack, stringval);
-	} else if (OPT_ISSET(ops,'v')) {
-	    setsparam(OPT_ARG(ops, 'v'), stringval);
+	if (READ_MSTREAM(buf,rcount,fout) < 0) {
+	    zwarnnam(name, "i/o error: %e", errno);
+	    if (buf)
+		free(buf);
 	} else {
-	    ent = prepnexthistent();
-	    ent->node.nam = stringval;
-	    ent->stim = ent->ftim = time(NULL);
-	    ent->node.flags = 0;
-	    ent->words = (short *)NULL;
-	    addhistnode(histtab, ent->node.nam, ent);
+	    stringval = metafy(buf, rcount, META_REALLOC);
+	    if (OPT_ISSET(ops,'z')) {
+		zpushnode(bufstack, stringval);
+	    } else if (OPT_ISSET(ops,'v')) {
+		setsparam(OPT_ARG(ops, 'v'), stringval);
+	    } else {
+		ent = prepnexthistent();
+		ent->node.nam = stringval;
+		ent->stim = ent->ftim = time(NULL);
+		ent->node.flags = 0;
+		ent->words = (short *)NULL;
+		addhistnode(histtab, ent->node.nam, ent);
+	    }
 	}
 	unqueue_signals();
     }
 
-#ifdef HAVE_OPEN_MEMSTREAM
-    if (fout)
-#endif
+    if (!CLOSE_CLEANLY(fout))
     {
-	/* Testing EBADF special-cases >&- redirections */
-	if ((fout != stdout) ? (fclose(fout) != 0) :
-	    (fflush(fout) != 0 && errno != EBADF)) {
-	    zwarnnam(name, "write error: %e", errno);
-	    ret = 1;
-	}
-	if (buf)
-	    free(buf);
+	zwarnnam(name, "write error: %e", errno);
+	ret = 1;
     }
     return ret;
 }