internal/web: consistently remove OMIT lines from code blocks

OMIT lines may appear in code before or after the section intended to be
displayed in a code block, however in the case of playground code blocks
those lines are actually included. This change makes sure we always
filter out the OMIT lines when generating code blocks, and adds a test.

I noticed this when the playground snippets in the race detector blog
post (https://go.dev/blog/race-detector) gave this error:

  package play: build constraints exclude all Go files in /tmp/sandbox2684341108

This is because those programs begin with the lines `// +build OMIT`,
which were never intended to be executed.

Change-Id: I5b0320da139741c789126f0c9d3dabc7044cd350
Reviewed-on: https://go-review.googlesource.com/c/website/+/466775
Run-TryBot: Russ Cox <rsc@golang.org>
Reviewed-by: Ian Lance Taylor <iant@google.com>
Auto-Submit: Russ Cox <rsc@golang.org>
Reviewed-by: Russ Cox <rsc@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
diff --git a/internal/web/code.go b/internal/web/code.go
index 03c0de1..e8bad4c 100644
--- a/internal/web/code.go
+++ b/internal/web/code.go
@@ -99,6 +99,9 @@
 		if cfg.Line == -1 {
 			cfg.Line = 1
 		}
+		lines := strings.SplitAfter(text, "\n")
+		filterOmit(lines)
+		text = strings.Join(lines, "")
 	case 1:
 		var n int
 		before, text, after, n = s.oneLine(file, text, arg[0])
@@ -153,8 +156,12 @@
 	if !isInt {
 		line = match(file, 0, lines, pattern)
 	}
+	filterOmit(lines)
 	line--
-	return strings.Join(lines[:line], ""), lines[line], strings.Join(lines[line+1:], ""), line
+	return strings.Join(lines[:line], ""),
+		lines[line],
+		strings.Join(lines[line+1:], ""),
+		line
 }
 
 // multipleLines returns the text generated by a three-argument code invocation.
@@ -170,15 +177,12 @@
 	} else if line2 < line1 {
 		log.Panicf("lines out of order for %q: %d %d", file, line1, line2)
 	}
-	for k := line1 - 1; k < line2; k++ {
-		if strings.HasSuffix(lines[k], "OMIT\n") {
-			lines[k] = ""
-		}
-	}
+	filterOmit(lines)
 	line1--
 	return strings.Join(lines[:line1], ""),
 		strings.Join(lines[line1:line2], ""),
-		strings.Join(lines[line2:], ""), line1
+		strings.Join(lines[line2:], ""),
+		line1
 }
 
 // parseArg returns the integer or string value of the argument and tells which it is.
@@ -223,3 +227,12 @@
 	log.Panicf("unrecognized pattern: %q", pattern)
 	return 0
 }
+
+func filterOmit(lines []string) {
+	for i, s := range lines {
+		if strings.HasSuffix(s, "OMIT\n") {
+			lines[i] = ""
+		}
+	}
+
+}
diff --git a/internal/web/site_test.go b/internal/web/site_test.go
index 1f27d28..793038f 100644
--- a/internal/web/site_test.go
+++ b/internal/web/site_test.go
@@ -63,6 +63,97 @@
 	testServeBody(t, site, "/doc/test2", "<em>template</em>")
 }
 
+func TestCode(t *testing.T) {
+	site := NewSite(fstest.MapFS{
+		"site.tmpl": {Data: []byte(`{{.Content}}`)},
+		"doc/code.md": {Data: []byte(`
+# hi
+whole file
+{{code "_code/prog.go"}}
+one line
+{{code "_code/prog.go" "/func main/"}}
+multiple lines
+{{code "_code/prog.go" "/START/" "/END/"}}
+following lines
+{{code "_code/prog.go" "/START/" "$"}}
+play
+{{play "_code/prog.go" "/START/" "/END/"}}
+play with numbers
+{{play "_code/prog.go" "/START/" "/END/" 0}}
+`)},
+		"doc/_code/prog.go": {Data: []byte(`
+// +build OMIT
+
+package main
+
+// START OMIT
+func main() { fmt.Println("hi") }
+// END OMIT
+
+func foo() {}
+`)},
+	})
+
+	testServeBody(t, site, "/doc/code", `<h1 id="hi">hi</h1>
+<p>whole file</p>
+<div class="code">
+<pre>package main
+
+func main() { fmt.Println(&#34;hi&#34;) }
+
+func foo() {}
+</pre>
+</div>
+<p>one line</p>
+<div class="code">
+<pre>func main() { fmt.Println(&#34;hi&#34;) }
+</pre>
+</div>
+<p>multiple lines</p>
+<div class="code">
+<pre>func main() { fmt.Println(&#34;hi&#34;) }
+</pre>
+</div>
+<p>following lines</p>
+<div class="code">
+<pre>func main() { fmt.Println(&#34;hi&#34;) }
+
+func foo() {}
+</pre>
+</div>
+<p>play</p>
+<div class="playground">
+<pre style="display: none"><span>
+
+package main
+
+</span>
+</pre>
+<pre contenteditable="true" spellcheck="false">func main() { fmt.Println(&#34;hi&#34;) }
+</pre>
+<pre style="display: none"><span>
+func foo() {}
+</span>
+</pre>
+</div>
+<p>play with numbers</p>
+<div class="playground">
+<pre style="display: none"><span>
+
+package main
+
+</span>
+</pre>
+<pre contenteditable="true" spellcheck="false"><span class="number"> 5&nbsp;&nbsp;</span>func main() { fmt.Println(&#34;hi&#34;) }
+<span class="number"> 6&nbsp;&nbsp;</span>
+</pre>
+<pre style="display: none"><span>
+func foo() {}
+</span>
+</pre>
+</div>`)
+}
+
 func TestTypeScript(t *testing.T) {
 	exampleOut, err := os.ReadFile("testdata/example.js")
 	if err != nil {