sha3: reenable s390x assembly

Fixes golang/go#64897

Change-Id: I0c8c52d73a7d2df0f44fee36d407a87213f59bff
Reviewed-on: https://go-review.googlesource.com/c/crypto/+/554435
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Cherry Mui <cherryyz@google.com>
Reviewed-by: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Roland Shoemaker <roland@golang.org>
Auto-Submit: Filippo Valsorda <filippo@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/sha3/allocations_test.go b/sha3/allocations_test.go
index c925099..36de5d5 100644
--- a/sha3/allocations_test.go
+++ b/sha3/allocations_test.go
@@ -7,6 +7,7 @@
 package sha3_test
 
 import (
+	"runtime"
 	"testing"
 
 	"golang.org/x/crypto/sha3"
@@ -15,6 +16,13 @@
 var sink byte
 
 func TestAllocations(t *testing.T) {
+	want := 0.0
+
+	if runtime.GOARCH == "s390x" {
+		// On s390x the returned hash.Hash is conditional so it escapes.
+		want = 3.0
+	}
+
 	t.Run("New", func(t *testing.T) {
 		if allocs := testing.AllocsPerRun(10, func() {
 			h := sha3.New256()
@@ -23,7 +31,7 @@
 			out := make([]byte, 0, 32)
 			out = h.Sum(out)
 			sink ^= out[0]
-		}); allocs > 0 {
+		}); allocs > want {
 			t.Errorf("expected zero allocations, got %0.1f", allocs)
 		}
 	})
@@ -37,7 +45,7 @@
 			sink ^= out[0]
 			h.Read(out)
 			sink ^= out[0]
-		}); allocs > 0 {
+		}); allocs > want {
 			t.Errorf("expected zero allocations, got %0.1f", allocs)
 		}
 	})
@@ -46,7 +54,7 @@
 			b := []byte("ABC")
 			out := sha3.Sum256(b)
 			sink ^= out[0]
-		}); allocs > 0 {
+		}); allocs > want {
 			t.Errorf("expected zero allocations, got %0.1f", allocs)
 		}
 	})
diff --git a/sha3/hashes.go b/sha3/hashes.go
index 1e815c9..5eae6cb 100644
--- a/sha3/hashes.go
+++ b/sha3/hashes.go
@@ -16,27 +16,43 @@
 // Its generic security strength is 224 bits against preimage attacks,
 // and 112 bits against collision attacks.
 func New224() hash.Hash {
-	return &state{rate: 144, outputLen: 28, dsbyte: 0x06}
+	return new224()
 }
 
 // New256 creates a new SHA3-256 hash.
 // Its generic security strength is 256 bits against preimage attacks,
 // and 128 bits against collision attacks.
 func New256() hash.Hash {
-	return &state{rate: 136, outputLen: 32, dsbyte: 0x06}
+	return new256()
 }
 
 // New384 creates a new SHA3-384 hash.
 // Its generic security strength is 384 bits against preimage attacks,
 // and 192 bits against collision attacks.
 func New384() hash.Hash {
-	return &state{rate: 104, outputLen: 48, dsbyte: 0x06}
+	return new384()
 }
 
 // New512 creates a new SHA3-512 hash.
 // Its generic security strength is 512 bits against preimage attacks,
 // and 256 bits against collision attacks.
 func New512() hash.Hash {
+	return new512()
+}
+
+func new224Generic() *state {
+	return &state{rate: 144, outputLen: 28, dsbyte: 0x06}
+}
+
+func new256Generic() *state {
+	return &state{rate: 136, outputLen: 32, dsbyte: 0x06}
+}
+
+func new384Generic() *state {
+	return &state{rate: 104, outputLen: 48, dsbyte: 0x06}
+}
+
+func new512Generic() *state {
 	return &state{rate: 72, outputLen: 64, dsbyte: 0x06}
 }
 
diff --git a/sha3/hashes_noasm.go b/sha3/hashes_noasm.go
new file mode 100644
index 0000000..9d85fb6
--- /dev/null
+++ b/sha3/hashes_noasm.go
@@ -0,0 +1,23 @@
+// Copyright 2023 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.
+
+//go:build !gc || purego || !s390x
+
+package sha3
+
+func new224() *state {
+	return new224Generic()
+}
+
+func new256() *state {
+	return new256Generic()
+}
+
+func new384() *state {
+	return new384Generic()
+}
+
+func new512() *state {
+	return new512Generic()
+}
diff --git a/sha3/sha3_s390x.go b/sha3/sha3_s390x.go
index 26b728b..00d8034 100644
--- a/sha3/sha3_s390x.go
+++ b/sha3/sha3_s390x.go
@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-//go:build gc && !purego && ignore
+//go:build gc && !purego
 
 package sha3
 
@@ -248,56 +248,56 @@
 	return s.clone()
 }
 
-// new224Asm returns an assembly implementation of SHA3-224 if available,
-// otherwise it returns nil.
-func new224Asm() hash.Hash {
+// new224 returns an assembly implementation of SHA3-224 if available,
+// otherwise it returns a generic implementation.
+func new224() hash.Hash {
 	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_224)
 	}
-	return nil
+	return new224Generic()
 }
 
-// new256Asm returns an assembly implementation of SHA3-256 if available,
-// otherwise it returns nil.
-func new256Asm() hash.Hash {
+// new256 returns an assembly implementation of SHA3-256 if available,
+// otherwise it returns a generic implementation.
+func new256() hash.Hash {
 	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_256)
 	}
-	return nil
+	return new256Generic()
 }
 
-// new384Asm returns an assembly implementation of SHA3-384 if available,
-// otherwise it returns nil.
-func new384Asm() hash.Hash {
+// new384 returns an assembly implementation of SHA3-384 if available,
+// otherwise it returns a generic implementation.
+func new384() hash.Hash {
 	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_384)
 	}
-	return nil
+	return new384Generic()
 }
 
-// new512Asm returns an assembly implementation of SHA3-512 if available,
-// otherwise it returns nil.
-func new512Asm() hash.Hash {
+// new512 returns an assembly implementation of SHA3-512 if available,
+// otherwise it returns a generic implementation.
+func new512() hash.Hash {
 	if cpu.S390X.HasSHA3 {
 		return newAsmState(sha3_512)
 	}
-	return nil
+	return new512Generic()
 }
 
-// newShake128Asm returns an assembly implementation of SHAKE-128 if available,
-// otherwise it returns nil.
-func newShake128Asm() ShakeHash {
+// newShake128 returns an assembly implementation of SHAKE-128 if available,
+// otherwise it returns a generic implementation.
+func newShake128() ShakeHash {
 	if cpu.S390X.HasSHA3 {
 		return newAsmState(shake_128)
 	}
-	return nil
+	return newShake128Generic()
 }
 
-// newShake256Asm returns an assembly implementation of SHAKE-256 if available,
-// otherwise it returns nil.
-func newShake256Asm() ShakeHash {
+// newShake256 returns an assembly implementation of SHAKE-256 if available,
+// otherwise it returns a generic implementation.
+func newShake256() ShakeHash {
 	if cpu.S390X.HasSHA3 {
 		return newAsmState(shake_256)
 	}
-	return nil
+	return newShake256Generic()
 }
diff --git a/sha3/sha3_s390x.s b/sha3/sha3_s390x.s
index df51683..826b862 100644
--- a/sha3/sha3_s390x.s
+++ b/sha3/sha3_s390x.s
@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-//go:build gc && !purego && ignore
+//go:build gc && !purego
 
 #include "textflag.h"
 
diff --git a/sha3/shake.go b/sha3/shake.go
index a31bcf8..1ea9275 100644
--- a/sha3/shake.go
+++ b/sha3/shake.go
@@ -115,13 +115,21 @@
 // Its generic security strength is 128 bits against all attacks if at
 // least 32 bytes of its output are used.
 func NewShake128() ShakeHash {
-	return &state{rate: rate128, outputLen: 32, dsbyte: dsbyteShake}
+	return newShake128()
 }
 
 // NewShake256 creates a new SHAKE256 variable-output-length ShakeHash.
 // Its generic security strength is 256 bits against all attacks if
 // at least 64 bytes of its output are used.
 func NewShake256() ShakeHash {
+	return newShake256()
+}
+
+func newShake128Generic() *state {
+	return &state{rate: rate128, outputLen: 32, dsbyte: dsbyteShake}
+}
+
+func newShake256Generic() *state {
 	return &state{rate: rate256, outputLen: 64, dsbyte: dsbyteShake}
 }
 
diff --git a/sha3/shake_noasm.go b/sha3/shake_noasm.go
new file mode 100644
index 0000000..4276ba4
--- /dev/null
+++ b/sha3/shake_noasm.go
@@ -0,0 +1,15 @@
+// Copyright 2023 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.
+
+//go:build !gc || purego || !s390x
+
+package sha3
+
+func newShake128() *state {
+	return newShake128Generic()
+}
+
+func newShake256() *state {
+	return newShake256Generic()
+}