From e1555a36d0068cd0cbc58915a3e2ea252a04e5ff Mon Sep 17 00:00:00 2001
From: Alan Donovan <adonovan@google.com>
Date: Fri, 19 Jan 2024 12:15:15 -0500
Subject: [PATCH] go/ssa: add Function.DomPostorder

Plus a test.

Fixes golang/go#46941

Change-Id: I8aa495bb18359fdc6acae86876a89687c1ba1d13
Reviewed-on: https://go-review.googlesource.com/c/tools/+/557055
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Tim King <taking@google.com>
---
 go/ssa/dom.go      | 31 ++++++++++++++-----------
 go/ssa/dom_test.go | 56 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 13 deletions(-)
 create mode 100644 go/ssa/dom_test.go

diff --git a/go/ssa/dom.go b/go/ssa/dom.go
index 66a2f5e6ed3..02c1ae83ae3 100644
--- a/go/ssa/dom.go
+++ b/go/ssa/dom.go
@@ -40,20 +40,25 @@ func (b *BasicBlock) Dominates(c *BasicBlock) bool {
 	return b.dom.pre <= c.dom.pre && c.dom.post <= b.dom.post
 }
 
-type byDomPreorder []*BasicBlock
-
-func (a byDomPreorder) Len() int           { return len(a) }
-func (a byDomPreorder) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
-func (a byDomPreorder) Less(i, j int) bool { return a[i].dom.pre < a[j].dom.pre }
-
-// DomPreorder returns a new slice containing the blocks of f in
-// dominator tree preorder.
+// DomPreorder returns a new slice containing the blocks of f
+// in a preorder traversal of the dominator tree.
 func (f *Function) DomPreorder() []*BasicBlock {
-	n := len(f.Blocks)
-	order := make(byDomPreorder, n)
-	copy(order, f.Blocks)
-	sort.Sort(order)
-	return order
+	slice := append([]*BasicBlock(nil), f.Blocks...)
+	sort.Slice(slice, func(i, j int) bool {
+		return slice[i].dom.pre < slice[j].dom.pre
+	})
+	return slice
+}
+
+// DomPostorder returns a new slice containing the blocks of f
+// in a postorder traversal of the dominator tree.
+// (This is not the same as a postdominance order.)
+func (f *Function) DomPostorder() []*BasicBlock {
+	slice := append([]*BasicBlock(nil), f.Blocks...)
+	sort.Slice(slice, func(i, j int) bool {
+		return slice[i].dom.post < slice[j].dom.post
+	})
+	return slice
 }
 
 // domInfo contains a BasicBlock's dominance information.
diff --git a/go/ssa/dom_test.go b/go/ssa/dom_test.go
new file mode 100644
index 00000000000..a98468fe101
--- /dev/null
+++ b/go/ssa/dom_test.go
@@ -0,0 +1,56 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package ssa_test
+
+import (
+	"fmt"
+	"path/filepath"
+	"testing"
+
+	"golang.org/x/tools/go/packages"
+	"golang.org/x/tools/go/ssa/ssautil"
+)
+
+func TestDominatorOrder(t *testing.T) {
+	const src = `package p
+
+func f(cond bool) {
+	// (Print operands match BasicBlock IDs.)
+	print(0)
+	if cond {
+		print(1)
+	} else {
+		print(2)
+	}
+	print(3)
+}
+`
+	dir := t.TempDir()
+	cfg := &packages.Config{
+		Dir:  dir,
+		Mode: packages.LoadSyntax,
+		Overlay: map[string][]byte{
+			filepath.Join(dir, "p.go"): []byte(src),
+		},
+	}
+	initial, err := packages.Load(cfg, "./p.go")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if packages.PrintErrors(initial) > 0 {
+		t.Fatal("packages contain errors")
+	}
+	_, pkgs := ssautil.Packages(initial, 0)
+	p := pkgs[0]
+	p.Build()
+	f := p.Func("f")
+
+	if got, want := fmt.Sprint(f.DomPreorder()), "[0 1 2 3]"; got != want {
+		t.Errorf("DomPreorder: got %v, want %s", got, want)
+	}
+	if got, want := fmt.Sprint(f.DomPostorder()), "[1 2 3 0]"; got != want {
+		t.Errorf("DomPostorder: got %v, want %s", got, want)
+	}
+}