3 files changed, 68 insertions, 15 deletions

diff --git a/src/cmd/compile/internal/ssa/check.go b/src/cmd/compile/internal/ssa/check.go
index d78e915091..d6d39aee76 100644
--- a/src/cmd/compile/internal/ssa/check.go
+++ b/src/cmd/compile/internal/ssa/check.go
@@ -294,6 +294,39 @@ func checkFunc(f *Func) {
 			}
 		}
 	}
+
+	// Check that if a tuple has a memory type, it is second.
+	for _, b := range f.Blocks {
+		for _, v := range b.Values {
+			if v.Type.IsTuple() && v.Type.FieldType(0).IsMemory() {
+				f.Fatalf("memory is first in a tuple: %s\n", v.LongString())
+			}
+		}
+	}
+
+	// Check that only one memory is live at any point.
+	// TODO: make this check examine interblock.
+	if f.scheduled {
+		for _, b := range f.Blocks {
+			var mem *Value // the live memory
+			for _, v := range b.Values {
+				if v.Op != OpPhi {
+					for _, a := range v.Args {
+						if a.Type.IsMemory() || a.Type.IsTuple() && a.Type.FieldType(1).IsMemory() {
+							if mem == nil {
+								mem = a
+							} else if mem != a {
+								f.Fatalf("two live mems @ %s: %s and %s", v, mem, a)
+							}
+						}
+					}
+				}
+				if v.Type.IsMemory() || v.Type.IsTuple() && v.Type.FieldType(1).IsMemory() {
+					mem = v
+				}
+			}
+		}
+	}
 }
 
 // domCheck reports whether x dominates y (including x==y).
diff --git a/src/cmd/compile/internal/ssa/schedule.go b/src/cmd/compile/internal/ssa/schedule.go
index a455a9a399..78b61f0959 100644
--- a/src/cmd/compile/internal/ssa/schedule.go
+++ b/src/cmd/compile/internal/ssa/schedule.go
@@ -148,19 +148,20 @@ func schedule(f *Func) {
 		}
 	}
 
+	// TODO: make this logic permanent in types.IsMemory?
+	isMem := func(v *Value) bool {
+		return v.Type.IsMemory() || v.Type.IsTuple() && v.Type.FieldType(1).IsMemory()
+	}
+
 	for _, b := range f.Blocks {
 		// Find store chain for block.
 		// Store chains for different blocks overwrite each other, so
 		// the calculated store chain is good only for this block.
 		for _, v := range b.Values {
-			if v.Op != OpPhi && v.Type.IsMemory() {
-				mem := v
-				if v.Op == OpSelect1 {
-					v = v.Args[0]
-				}
+			if v.Op != OpPhi && isMem(v) {
 				for _, w := range v.Args {
-					if w.Type.IsMemory() {
-						nextMem[w.ID] = mem
+					if isMem(w) {
+						nextMem[w.ID] = v
 					}
 				}
 			}
@@ -179,15 +180,15 @@ func schedule(f *Func) {
 					uses[w.ID]++
 				}
 				// Any load must come before the following store.
-				if v.Type.IsMemory() || !w.Type.IsMemory() {
-					continue // not a load
-				}
-				s := nextMem[w.ID]
-				if s == nil || s.Block != b {
-					continue
+				if !isMem(v) && isMem(w) {
+					// v is a load.
+					s := nextMem[w.ID]
+					if s == nil || s.Block != b {
+						continue
+					}
+					additionalArgs[s.ID] = append(additionalArgs[s.ID], v)
+					uses[v.ID]++
 				}
-				additionalArgs[s.ID] = append(additionalArgs[s.ID], v)
-				uses[v.ID]++
 			}
 		}
 
diff --git a/test/fixedbugs/issue20335.go b/test/fixedbugs/issue20335.go
new file mode 100644
index 0000000000..185c2f06ea
--- /dev/null
+++ b/test/fixedbugs/issue20335.go
@@ -0,0 +1,19 @@
+// compile
+
+// Copyright 2017 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.
+
+// Issue 20335: don't reorder loads with stores.
+// This test should fail on the ssacheck builder
+// without the fix in the CL that added this file.
+// TODO: check the generated assembly?
+
+package a
+
+import "sync/atomic"
+
+func f(p, q *int32) bool {
+	x := *q
+	return atomic.AddInt32(p, 1) == x
+}