poly1305: rewrite the Go implementation with 64-bit limbs

The new code is meant to be readable without external references for
Poly1305, and explains the field logic. The generic code is now 30-50%
faster on a Intel(R) Xeon(R) CPU E5-2690 v3 @ 2.60GHz, and even better
on a 3.1 GHz i7 MacBook.

name        old time/op    new time/op    delta
64-48          126ns ± 0%      80ns ± 1%  -36.24%  (p=0.000 n=16+20)
1K-48         1.07µs ± 0%    0.81µs ± 2%  -23.63%  (p=0.000 n=19+20)
2M-48         2.07ms ± 0%    1.61ms ± 1%  -22.31%  (p=0.000 n=20+20)
Write64-48    79.3ns ± 0%    58.0ns ± 1%  -26.89%  (p=0.000 n=20+19)
Write1K-48    1.02µs ± 0%    0.79µs ± 1%  -22.91%  (p=0.000 n=19+19)
Write2M-48    2.07ms ± 0%    1.61ms ± 2%  -22.33%  (p=0.000 n=17+20)

name        old speed      new speed      delta
64-48        508MB/s ± 0%   797MB/s ± 1%  +56.95%  (p=0.000 n=16+20)
1K-48        960MB/s ± 0%  1257MB/s ± 2%  +30.94%  (p=0.000 n=18+20)
2M-48       1.01GB/s ± 0%  1.30GB/s ± 1%  +28.73%  (p=0.000 n=20+20)
Write64-48   807MB/s ± 0%  1104MB/s ± 1%  +36.78%  (p=0.000 n=18+19)
Write1K-48  1.00GB/s ± 0%  1.30GB/s ± 1%  +29.71%  (p=0.000 n=18+19)
Write2M-48  1.01GB/s ± 0%  1.31GB/s ± 2%  +28.77%  (p=0.000 n=17+20)

The assembly is still 50-90% faster on the Xeon, 30-60% on the MacBook.
The Go code does not use all the arithmetic tricks the assembly does,
and it does not have access to the three operand wide shift instruction.

name        old time/op    new time/op    delta
64-48         80.3ns ± 1%    54.2ns ± 0%  -32.50%  (p=0.000 n=20+17)
1K-48          815ns ± 2%     446ns ± 1%  -45.27%  (p=0.000 n=20+20)
2M-48         1.61ms ± 1%    0.86ms ± 0%  -46.54%  (p=0.000 n=20+17)
Write64-48    58.0ns ± 1%    34.0ns ± 0%  -41.34%  (p=0.000 n=19+20)
Write1K-48     790ns ± 1%     427ns ± 0%  -45.92%  (p=0.000 n=19+17)
Write2M-48    1.61ms ± 2%    0.86ms ± 0%  -46.51%  (p=0.000 n=20+20)

name        old speed      new speed      delta
64-48        797MB/s ± 1%  1180MB/s ± 0%  +48.09%  (p=0.000 n=20+19)
1K-48       1.26GB/s ± 2%  2.30GB/s ± 1%  +82.71%  (p=0.000 n=20+20)
2M-48       1.30GB/s ± 1%  2.44GB/s ± 0%  +87.04%  (p=0.000 n=20+17)
Write64-48  1.10GB/s ± 1%  1.88GB/s ± 0%  +70.52%  (p=0.000 n=19+18)
Write1K-48  1.30GB/s ± 1%  2.40GB/s ± 0%  +84.84%  (p=0.000 n=19+18)
Write2M-48  1.31GB/s ± 2%  2.44GB/s ± 0%  +86.93%  (p=0.000 n=20+20)

Hopefully this will also avoid the need for an arm64 implementation.

Since now the Go and the amd64/ppc64le assembly use the same limb
schedule, drop the assembly initialize and finalize implementations,
and make the wrapper code match. It comes with a minor slowdown.

name        old time/op    new time/op    delta
64-48         50.3ns ± 0%    54.2ns ± 0%  +7.73%  (p=0.000 n=20+17)
1K-48          441ns ± 0%     446ns ± 1%  +1.10%  (p=0.000 n=19+20)
2M-48          860µs ± 0%     859µs ± 0%    ~     (p=0.178 n=19+17)
Write64-48    34.0ns ± 0%    34.0ns ± 0%    ~     (all equal)
Write1K-48     424ns ± 0%     427ns ± 0%  +0.71%  (p=0.000 n=17+17)
Write2M-48     860µs ± 0%     859µs ± 0%  -0.04%  (p=0.000 n=19+20)

name        old speed      new speed      delta
64-48       1.27GB/s ± 0%  1.18GB/s ± 0%  -7.20%  (p=0.000 n=20+19)
1K-48       2.32GB/s ± 0%  2.30GB/s ± 1%  -1.07%  (p=0.000 n=18+20)
2M-48       2.44GB/s ± 0%  2.44GB/s ± 0%    ~     (p=0.173 n=19+17)
Write64-48  1.88GB/s ± 0%  1.88GB/s ± 0%  +0.04%  (p=0.000 n=19+18)
Write1K-48  2.41GB/s ± 0%  2.40GB/s ± 0%  -0.67%  (p=0.000 n=19+18)
Write2M-48  2.44GB/s ± 0%  2.44GB/s ± 0%  +0.04%  (p=0.000 n=19+20)

Since poly1305/sum_generic.go was almost entirely rewritten, it's
probably best reviewed on gitiles.

This is the implementation published at
https://blog.filippo.io/a-literate-go-implementation-of-poly1305/

Updates #31470

Change-Id: I74f9011d3ee317a43b05ae7f05d96081d08bffd3
Reviewed-on: https://go-review.googlesource.com/c/crypto/+/169037
Reviewed-by: Katie Hockman <katie@golang.org>
diff --git a/poly1305/bits_compat.go b/poly1305/bits_compat.go
new file mode 100644
index 0000000..157a69f
--- /dev/null
+++ b/poly1305/bits_compat.go
@@ -0,0 +1,39 @@
+// Copyright 2019 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.
+
+// +build !go1.13
+
+package poly1305
+
+// Generic fallbacks for the math/bits intrinsics, copied from
+// src/math/bits/bits.go. They were added in Go 1.12, but Add64 and Sum64 had
+// variable time fallbacks until Go 1.13.
+
+func bitsAdd64(x, y, carry uint64) (sum, carryOut uint64) {
+	sum = x + y + carry
+	carryOut = ((x & y) | ((x | y) &^ sum)) >> 63
+	return
+}
+
+func bitsSub64(x, y, borrow uint64) (diff, borrowOut uint64) {
+	diff = x - y - borrow
+	borrowOut = ((^x & y) | (^(x ^ y) & diff)) >> 63
+	return
+}
+
+func bitsMul64(x, y uint64) (hi, lo uint64) {
+	const mask32 = 1<<32 - 1
+	x0 := x & mask32
+	x1 := x >> 32
+	y0 := y & mask32
+	y1 := y >> 32
+	w0 := x0 * y0
+	t := x1*y0 + w0>>32
+	w1 := t & mask32
+	w2 := t >> 32
+	w1 += x0 * y1
+	hi = x1*y1 + w2 + w1>>32
+	lo = x * y
+	return
+}
diff --git a/poly1305/bits_go1.13.go b/poly1305/bits_go1.13.go
new file mode 100644
index 0000000..a0a185f
--- /dev/null
+++ b/poly1305/bits_go1.13.go
@@ -0,0 +1,21 @@
+// Copyright 2019 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.
+
+// +build go1.13
+
+package poly1305
+
+import "math/bits"
+
+func bitsAdd64(x, y, carry uint64) (sum, carryOut uint64) {
+	return bits.Add64(x, y, carry)
+}
+
+func bitsSub64(x, y, borrow uint64) (diff, borrowOut uint64) {
+	return bits.Sub64(x, y, borrow)
+}
+
+func bitsMul64(x, y uint64) (hi, lo uint64) {
+	return bits.Mul64(x, y)
+}
diff --git a/poly1305/poly1305.go b/poly1305/poly1305.go
index d076a56..066159b 100644
--- a/poly1305/poly1305.go
+++ b/poly1305/poly1305.go
@@ -22,8 +22,14 @@
 // TagSize is the size, in bytes, of a poly1305 authenticator.
 const TagSize = 16
 
-// Verify returns true if mac is a valid authenticator for m with the given
-// key.
+// Sum generates an authenticator for msg using a one-time key and puts the
+// 16-byte result into out. Authenticating two different messages with the same
+// key allows an attacker to forge messages at will.
+func Sum(out *[16]byte, m []byte, key *[32]byte) {
+	sum(out, m, key)
+}
+
+// Verify returns true if mac is a valid authenticator for m with the given key.
 func Verify(mac *[16]byte, m []byte, key *[32]byte) bool {
 	var tmp [16]byte
 	Sum(&tmp, m, key)
diff --git a/poly1305/poly1305_test.go b/poly1305/poly1305_test.go
index e35928a..b258eed 100644
--- a/poly1305/poly1305_test.go
+++ b/poly1305/poly1305_test.go
@@ -5,6 +5,7 @@
 package poly1305
 
 import (
+	"crypto/rand"
 	"encoding/hex"
 	"flag"
 	"testing"
@@ -115,8 +116,14 @@
 			input = unalignBytes(input)
 		}
 		h := newMACGeneric(&key)
-		h.Write(input[:len(input)/2])
-		h.Write(input[len(input)/2:])
+		n, err := h.Write(input[:len(input)/3])
+		if err != nil || n != len(input[:len(input)/3]) {
+			t.Errorf("#%d: unexpected Write results: n = %d, err = %v", i, n, err)
+		}
+		n, err = h.Write(input[len(input)/3:])
+		if err != nil || n != len(input[len(input)/3:]) {
+			t.Errorf("#%d: unexpected Write results: n = %d, err = %v", i, n, err)
+		}
 		h.Sum(&out)
 		if tag := v.Tag(); out != tag {
 			t.Errorf("%d: expected %x, got %x", i, tag[:], out[:])
@@ -134,8 +141,14 @@
 			input = unalignBytes(input)
 		}
 		h := New(&key)
-		h.Write(input[:len(input)/2])
-		h.Write(input[len(input)/2:])
+		n, err := h.Write(input[:len(input)/3])
+		if err != nil || n != len(input[:len(input)/3]) {
+			t.Errorf("#%d: unexpected Write results: n = %d, err = %v", i, n, err)
+		}
+		n, err = h.Write(input[len(input)/3:])
+		if err != nil || n != len(input[len(input)/3:]) {
+			t.Errorf("#%d: unexpected Write results: n = %d, err = %v", i, n, err)
+		}
 		h.Sum(out[:0])
 		if tag := v.Tag(); out != tag {
 			t.Errorf("%d: expected %x, got %x", i, tag[:], out[:])
@@ -150,6 +163,7 @@
 	if unaligned {
 		in = unalignBytes(in)
 	}
+	rand.Read(in)
 	b.SetBytes(int64(len(in)))
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
@@ -164,6 +178,7 @@
 	if unaligned {
 		in = unalignBytes(in)
 	}
+	rand.Read(in)
 	b.SetBytes(int64(len(in)))
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
diff --git a/poly1305/sum_amd64.go b/poly1305/sum_amd64.go
index 2dbf42a..df56a65 100644
--- a/poly1305/sum_amd64.go
+++ b/poly1305/sum_amd64.go
@@ -7,62 +7,52 @@
 package poly1305
 
 //go:noescape
-func initialize(state *[7]uint64, key *[32]byte)
+func update(state *macState, msg []byte)
 
-//go:noescape
-func update(state *[7]uint64, msg []byte)
-
-//go:noescape
-func finalize(tag *[TagSize]byte, state *[7]uint64)
-
-// Sum generates an authenticator for m using a one-time key and puts the
-// 16-byte result into out. Authenticating two different messages with the same
-// key allows an attacker to forge messages at will.
-func Sum(out *[16]byte, m []byte, key *[32]byte) {
+func sum(out *[16]byte, m []byte, key *[32]byte) {
 	h := newMAC(key)
 	h.Write(m)
 	h.Sum(out)
 }
 
 func newMAC(key *[32]byte) (h mac) {
-	initialize(&h.state, key)
+	initialize(key, &h.r, &h.s)
 	return
 }
 
-type mac struct {
-	state [7]uint64 // := uint64{ h0, h1, h2, r0, r1, pad0, pad1 }
+// mac is a wrapper for macGeneric that redirects calls that would have gone to
+// updateGeneric to update.
+//
+// Its Write and Sum methods are otherwise identical to the macGeneric ones, but
+// using function pointers would carry a major performance cost.
+type mac struct{ macGeneric }
 
-	buffer [TagSize]byte
-	offset int
-}
-
-func (h *mac) Write(p []byte) (n int, err error) {
-	n = len(p)
+func (h *mac) Write(p []byte) (int, error) {
+	nn := len(p)
 	if h.offset > 0 {
-		remaining := TagSize - h.offset
-		if n < remaining {
-			h.offset += copy(h.buffer[h.offset:], p)
-			return n, nil
+		n := copy(h.buffer[h.offset:], p)
+		if h.offset+n < TagSize {
+			h.offset += n
+			return nn, nil
 		}
-		copy(h.buffer[h.offset:], p[:remaining])
-		p = p[remaining:]
+		p = p[n:]
 		h.offset = 0
-		update(&h.state, h.buffer[:])
+		update(&h.macState, h.buffer[:])
 	}
-	if nn := len(p) - (len(p) % TagSize); nn > 0 {
-		update(&h.state, p[:nn])
-		p = p[nn:]
+	if n := len(p) - (len(p) % TagSize); n > 0 {
+		update(&h.macState, p[:n])
+		p = p[n:]
 	}
 	if len(p) > 0 {
 		h.offset += copy(h.buffer[h.offset:], p)
 	}
-	return n, nil
+	return nn, nil
 }
 
 func (h *mac) Sum(out *[16]byte) {
-	state := h.state
+	state := h.macState
 	if h.offset > 0 {
 		update(&state, h.buffer[:h.offset])
 	}
-	finalize(out, &state)
+	finalize(out, &state.h, &state.s)
 }
diff --git a/poly1305/sum_amd64.s b/poly1305/sum_amd64.s
index 7d600f1..8c0cefb 100644
--- a/poly1305/sum_amd64.s
+++ b/poly1305/sum_amd64.s
@@ -54,10 +54,6 @@
 	ADCQ  t3, h1;                  \
 	ADCQ  $0, h2
 
-DATA ·poly1305Mask<>+0x00(SB)/8, $0x0FFFFFFC0FFFFFFF
-DATA ·poly1305Mask<>+0x08(SB)/8, $0x0FFFFFFC0FFFFFFC
-GLOBL ·poly1305Mask<>(SB), RODATA, $16
-
 // func update(state *[7]uint64, msg []byte)
 TEXT ·update(SB), $0-32
 	MOVQ state+0(FP), DI
@@ -110,39 +106,3 @@
 	MOVQ R9, 8(DI)
 	MOVQ R10, 16(DI)
 	RET
-
-// func initialize(state *[7]uint64, key *[32]byte)
-TEXT ·initialize(SB), $0-16
-	MOVQ state+0(FP), DI
-	MOVQ key+8(FP), SI
-
-	// state[0...7] is initialized with zero
-	MOVOU 0(SI), X0
-	MOVOU 16(SI), X1
-	MOVOU ·poly1305Mask<>(SB), X2
-	PAND  X2, X0
-	MOVOU X0, 24(DI)
-	MOVOU X1, 40(DI)
-	RET
-
-// func finalize(tag *[TagSize]byte, state *[7]uint64)
-TEXT ·finalize(SB), $0-16
-	MOVQ tag+0(FP), DI
-	MOVQ state+8(FP), SI
-
-	MOVQ    0(SI), AX
-	MOVQ    8(SI), BX
-	MOVQ    16(SI), CX
-	MOVQ    AX, R8
-	MOVQ    BX, R9
-	SUBQ    $0xFFFFFFFFFFFFFFFB, AX
-	SBBQ    $0xFFFFFFFFFFFFFFFF, BX
-	SBBQ    $3, CX
-	CMOVQCS R8, AX
-	CMOVQCS R9, BX
-	ADDQ    40(SI), AX
-	ADCQ    48(SI), BX
-
-	MOVQ AX, 0(DI)
-	MOVQ BX, 8(DI)
-	RET
diff --git a/poly1305/sum_arm.go b/poly1305/sum_arm.go
index 5dc321c..6e695e4 100644
--- a/poly1305/sum_arm.go
+++ b/poly1305/sum_arm.go
@@ -6,14 +6,11 @@
 
 package poly1305
 
-// This function is implemented in sum_arm.s
+// poly1305_auth_armv6 is implemented in sum_arm.s
 //go:noescape
 func poly1305_auth_armv6(out *[16]byte, m *byte, mlen uint32, key *[32]byte)
 
-// Sum generates an authenticator for m using a one-time key and puts the
-// 16-byte result into out. Authenticating two different messages with the same
-// key allows an attacker to forge messages at will.
-func Sum(out *[16]byte, m []byte, key *[32]byte) {
+func sum(out *[16]byte, m []byte, key *[32]byte) {
 	var mPtr *byte
 	if len(m) > 0 {
 		mPtr = &m[0]
diff --git a/poly1305/sum_generic.go b/poly1305/sum_generic.go
index bab76ef..1187eab 100644
--- a/poly1305/sum_generic.go
+++ b/poly1305/sum_generic.go
@@ -2,18 +2,29 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+// This file provides the generic implementation of Sum and MAC. Other files
+// might provide optimized assembly implementations of some of this code.
+
 package poly1305
 
 import "encoding/binary"
 
-const (
-	msgBlock   = uint32(1 << 24)
-	finalBlock = uint32(0)
-)
+// Poly1305 [RFC 7539] is a relatively simple algorithm: the authentication tag
+// for a 64 bytes message is approximately
+//
+//     s + m[0:16] * r⁴ + m[16:32] * r³ + m[32:48] * r² + m[48:64] * r  mod  2¹³⁰ - 5
+//
+// for some secret r and s. It can be computed sequentially like
+//
+//     for len(msg) > 0:
+//         h += read(msg, 16)
+//         h *= r
+//         h %= 2¹³⁰ - 5
+//     return h + s
+//
+// All the complexity is about doing performant constant-time math on numbers
+// larger than any available numeric type.
 
-// sumGeneric generates an authenticator for msg using a one-time key and
-// puts the 16-byte result into out. This is the generic implementation of
-// Sum and should be called if no assembly implementation is available.
 func sumGeneric(out *[TagSize]byte, msg []byte, key *[32]byte) {
 	h := newMACGeneric(key)
 	h.Write(msg)
@@ -21,152 +32,276 @@
 }
 
 func newMACGeneric(key *[32]byte) (h macGeneric) {
-	h.r[0] = binary.LittleEndian.Uint32(key[0:]) & 0x3ffffff
-	h.r[1] = (binary.LittleEndian.Uint32(key[3:]) >> 2) & 0x3ffff03
-	h.r[2] = (binary.LittleEndian.Uint32(key[6:]) >> 4) & 0x3ffc0ff
-	h.r[3] = (binary.LittleEndian.Uint32(key[9:]) >> 6) & 0x3f03fff
-	h.r[4] = (binary.LittleEndian.Uint32(key[12:]) >> 8) & 0x00fffff
-
-	h.s[0] = binary.LittleEndian.Uint32(key[16:])
-	h.s[1] = binary.LittleEndian.Uint32(key[20:])
-	h.s[2] = binary.LittleEndian.Uint32(key[24:])
-	h.s[3] = binary.LittleEndian.Uint32(key[28:])
+	initialize(key, &h.r, &h.s)
 	return
 }
 
+// macState holds numbers in saturated 64-bit little-endian limbs. That is,
+// the value of [x0, x1, x2] is x[0] + x[1] * 2⁶⁴ + x[2] * 2¹²⁸.
+type macState struct {
+	// h is the main accumulator. It is to be interpreted modulo 2¹³⁰ - 5, but
+	// can grow larger during and after rounds.
+	h [3]uint64
+	// r and s are the private key components.
+	r [2]uint64
+	s [2]uint64
+}
+
 type macGeneric struct {
-	h, r [5]uint32
-	s    [4]uint32
+	macState
 
 	buffer [TagSize]byte
 	offset int
 }
 
-func (h *macGeneric) Write(p []byte) (n int, err error) {
-	n = len(p)
+// Write splits the incoming message into TagSize chunks, and passes them to
+// update. It buffers incomplete chunks.
+func (h *macGeneric) Write(p []byte) (int, error) {
+	nn := len(p)
 	if h.offset > 0 {
-		remaining := TagSize - h.offset
-		if n < remaining {
-			h.offset += copy(h.buffer[h.offset:], p)
-			return n, nil
+		n := copy(h.buffer[h.offset:], p)
+		if h.offset+n < TagSize {
+			h.offset += n
+			return nn, nil
 		}
-		copy(h.buffer[h.offset:], p[:remaining])
-		p = p[remaining:]
+		p = p[n:]
 		h.offset = 0
-		updateGeneric(h.buffer[:], msgBlock, &(h.h), &(h.r))
+		updateGeneric(&h.macState, h.buffer[:])
 	}
-	if nn := len(p) - (len(p) % TagSize); nn > 0 {
-		updateGeneric(p, msgBlock, &(h.h), &(h.r))
-		p = p[nn:]
+	if n := len(p) - (len(p) % TagSize); n > 0 {
+		updateGeneric(&h.macState, p[:n])
+		p = p[n:]
 	}
 	if len(p) > 0 {
 		h.offset += copy(h.buffer[h.offset:], p)
 	}
-	return n, nil
+	return nn, nil
 }
 
-func (h *macGeneric) Sum(out *[16]byte) {
-	H, R := h.h, h.r
+// Sum flushes the last incomplete chunk from the buffer, if any, and generates
+// the MAC output. It does not modify its state, in order to allow for multiple
+// calls to Sum, even if no Write is allowed after Sum.
+func (h *macGeneric) Sum(out *[TagSize]byte) {
+	state := h.macState
 	if h.offset > 0 {
-		var buffer [TagSize]byte
-		copy(buffer[:], h.buffer[:h.offset])
-		buffer[h.offset] = 1 // invariant: h.offset < TagSize
-		updateGeneric(buffer[:], finalBlock, &H, &R)
+		updateGeneric(&state, h.buffer[:h.offset])
 	}
-	finalizeGeneric(out, &H, &(h.s))
+	finalize(out, &state.h, &state.s)
 }
 
-func updateGeneric(msg []byte, flag uint32, h, r *[5]uint32) {
-	h0, h1, h2, h3, h4 := h[0], h[1], h[2], h[3], h[4]
-	r0, r1, r2, r3, r4 := uint64(r[0]), uint64(r[1]), uint64(r[2]), uint64(r[3]), uint64(r[4])
-	R1, R2, R3, R4 := r1*5, r2*5, r3*5, r4*5
+// [rMask0, rMask1] is the specified Poly1305 clamping mask in little-endian. It
+// clears some bits of the secret coefficient to make it possible to implement
+// multiplication more efficiently.
+const (
+	rMask0 = 0x0FFFFFFC0FFFFFFF
+	rMask1 = 0x0FFFFFFC0FFFFFFC
+)
 
-	for len(msg) >= TagSize {
-		// h += msg
-		h0 += binary.LittleEndian.Uint32(msg[0:]) & 0x3ffffff
-		h1 += (binary.LittleEndian.Uint32(msg[3:]) >> 2) & 0x3ffffff
-		h2 += (binary.LittleEndian.Uint32(msg[6:]) >> 4) & 0x3ffffff
-		h3 += (binary.LittleEndian.Uint32(msg[9:]) >> 6) & 0x3ffffff
-		h4 += (binary.LittleEndian.Uint32(msg[12:]) >> 8) | flag
-
-		// h *= r
-		d0 := (uint64(h0) * r0) + (uint64(h1) * R4) + (uint64(h2) * R3) + (uint64(h3) * R2) + (uint64(h4) * R1)
-		d1 := (d0 >> 26) + (uint64(h0) * r1) + (uint64(h1) * r0) + (uint64(h2) * R4) + (uint64(h3) * R3) + (uint64(h4) * R2)
-		d2 := (d1 >> 26) + (uint64(h0) * r2) + (uint64(h1) * r1) + (uint64(h2) * r0) + (uint64(h3) * R4) + (uint64(h4) * R3)
-		d3 := (d2 >> 26) + (uint64(h0) * r3) + (uint64(h1) * r2) + (uint64(h2) * r1) + (uint64(h3) * r0) + (uint64(h4) * R4)
-		d4 := (d3 >> 26) + (uint64(h0) * r4) + (uint64(h1) * r3) + (uint64(h2) * r2) + (uint64(h3) * r1) + (uint64(h4) * r0)
-
-		// h %= p
-		h0 = uint32(d0) & 0x3ffffff
-		h1 = uint32(d1) & 0x3ffffff
-		h2 = uint32(d2) & 0x3ffffff
-		h3 = uint32(d3) & 0x3ffffff
-		h4 = uint32(d4) & 0x3ffffff
-
-		h0 += uint32(d4>>26) * 5
-		h1 += h0 >> 26
-		h0 = h0 & 0x3ffffff
-
-		msg = msg[TagSize:]
-	}
-
-	h[0], h[1], h[2], h[3], h[4] = h0, h1, h2, h3, h4
+func initialize(key *[32]byte, r, s *[2]uint64) {
+	r[0] = binary.LittleEndian.Uint64(key[0:8]) & rMask0
+	r[1] = binary.LittleEndian.Uint64(key[8:16]) & rMask1
+	s[0] = binary.LittleEndian.Uint64(key[16:24])
+	s[1] = binary.LittleEndian.Uint64(key[24:32])
 }
 
-func finalizeGeneric(out *[TagSize]byte, h *[5]uint32, s *[4]uint32) {
-	h0, h1, h2, h3, h4 := h[0], h[1], h[2], h[3], h[4]
+// uint128 holds a 128-bit number as two 64-bit limbs, for use with the
+// bits.Mul64 and bits.Add64 intrinsics.
+type uint128 struct {
+	lo, hi uint64
+}
 
-	// h %= p reduction
-	h2 += h1 >> 26
-	h1 &= 0x3ffffff
-	h3 += h2 >> 26
-	h2 &= 0x3ffffff
-	h4 += h3 >> 26
-	h3 &= 0x3ffffff
-	h0 += 5 * (h4 >> 26)
-	h4 &= 0x3ffffff
-	h1 += h0 >> 26
-	h0 &= 0x3ffffff
+func mul64(a, b uint64) uint128 {
+	hi, lo := bitsMul64(a, b)
+	return uint128{lo, hi}
+}
 
-	// h - p
-	t0 := h0 + 5
-	t1 := h1 + (t0 >> 26)
-	t2 := h2 + (t1 >> 26)
-	t3 := h3 + (t2 >> 26)
-	t4 := h4 + (t3 >> 26) - (1 << 26)
-	t0 &= 0x3ffffff
-	t1 &= 0x3ffffff
-	t2 &= 0x3ffffff
-	t3 &= 0x3ffffff
+func add128(a, b uint128) uint128 {
+	lo, c := bitsAdd64(a.lo, b.lo, 0)
+	hi, c := bitsAdd64(a.hi, b.hi, c)
+	if c != 0 {
+		panic("poly1305: unexpected overflow")
+	}
+	return uint128{lo, hi}
+}
 
-	// select h if h < p else h - p
-	t_mask := (t4 >> 31) - 1
-	h_mask := ^t_mask
-	h0 = (h0 & h_mask) | (t0 & t_mask)
-	h1 = (h1 & h_mask) | (t1 & t_mask)
-	h2 = (h2 & h_mask) | (t2 & t_mask)
-	h3 = (h3 & h_mask) | (t3 & t_mask)
-	h4 = (h4 & h_mask) | (t4 & t_mask)
+func shiftRightBy2(a uint128) uint128 {
+	a.lo = a.lo>>2 | (a.hi&3)<<62
+	a.hi = a.hi >> 2
+	return a
+}
 
-	// h %= 2^128
-	h0 |= h1 << 26
-	h1 = ((h1 >> 6) | (h2 << 20))
-	h2 = ((h2 >> 12) | (h3 << 14))
-	h3 = ((h3 >> 18) | (h4 << 8))
+// updateGeneric absorbs msg into the state.h accumulator. For each chunk m of
+// 128 bits of message, it computes
+//
+//     h₊ = (h + m) * r  mod  2¹³⁰ - 5
+//
+// If the msg length is not a multiple of TagSize, it assumes the last
+// incomplete chunk is the final one.
+func updateGeneric(state *macState, msg []byte) {
+	h0, h1, h2 := state.h[0], state.h[1], state.h[2]
+	r0, r1 := state.r[0], state.r[1]
 
-	// s: the s part of the key
-	// tag = (h + s) % (2^128)
-	t := uint64(h0) + uint64(s[0])
-	h0 = uint32(t)
-	t = uint64(h1) + uint64(s[1]) + (t >> 32)
-	h1 = uint32(t)
-	t = uint64(h2) + uint64(s[2]) + (t >> 32)
-	h2 = uint32(t)
-	t = uint64(h3) + uint64(s[3]) + (t >> 32)
-	h3 = uint32(t)
+	for len(msg) > 0 {
+		var c uint64
 
-	binary.LittleEndian.PutUint32(out[0:], h0)
-	binary.LittleEndian.PutUint32(out[4:], h1)
-	binary.LittleEndian.PutUint32(out[8:], h2)
-	binary.LittleEndian.PutUint32(out[12:], h3)
+		// For the first step, h + m, we use a chain of bits.Add64 intrinsics.
+		// The resulting value of h might exceed 2¹³⁰ - 5, but will be partially
+		// reduced at the end of the multiplication below.
+		//
+		// The spec requires us to set a bit just above the message size, not to
+		// hide leading zeroes. For full chunks, that's 1 << 128, so we can just
+		// add 1 to the most significant (2¹²⁸) limb, h2.
+		if len(msg) >= TagSize {
+			h0, c = bitsAdd64(h0, binary.LittleEndian.Uint64(msg[0:8]), 0)
+			h1, c = bitsAdd64(h1, binary.LittleEndian.Uint64(msg[8:16]), c)
+			h2 += c + 1
+
+			msg = msg[TagSize:]
+		} else {
+			var buf [TagSize]byte
+			copy(buf[:], msg)
+			buf[len(msg)] = 1
+
+			h0, c = bitsAdd64(h0, binary.LittleEndian.Uint64(buf[0:8]), 0)
+			h1, c = bitsAdd64(h1, binary.LittleEndian.Uint64(buf[8:16]), c)
+			h2 += c
+
+			msg = nil
+		}
+
+		// Multiplication of big number limbs is similar to elementary school
+		// columnar multiplication. Instead of digits, there are 64-bit limbs.
+		//
+		// We are multiplying a 3 limbs number, h, by a 2 limbs number, r.
+		//
+		//                        h2    h1    h0  x
+		//                              r1    r0  =
+		//                       ----------------
+		//                      h2r0  h1r0  h0r0     <-- individual 128-bit products
+		//            +   h2r1  h1r1  h0r1
+		//               ------------------------
+		//                 m3    m2    m1    m0      <-- result in 128-bit overlapping limbs
+		//               ------------------------
+		//         m3.hi m2.hi m1.hi m0.hi           <-- carry propagation
+		//     +         m3.lo m2.lo m1.lo m0.lo
+		//        -------------------------------
+		//           t4    t3    t2    t1    t0      <-- final result in 64-bit limbs
+		//
+		// The main difference from pen-and-paper multiplication is that we do
+		// carry propagation in a separate step, as if we wrote two digit sums
+		// at first (the 128-bit limbs), and then carried the tens all at once.
+
+		h0r0 := mul64(h0, r0)
+		h1r0 := mul64(h1, r0)
+		h2r0 := mul64(h2, r0)
+		h0r1 := mul64(h0, r1)
+		h1r1 := mul64(h1, r1)
+		h2r1 := mul64(h2, r1)
+
+		// Since h2 is known to be at most 7 (5 + 1 + 1), and r0 and r1 have their
+		// top 4 bits cleared by rMask{0,1}, we know that their product is not going
+		// to overflow 64 bits, so we can ignore the high part of the products.
+		//
+		// This also means that the product doesn't have a fifth limb (t4).
+		if h2r0.hi != 0 {
+			panic("poly1305: unexpected overflow")
+		}
+		if h2r1.hi != 0 {
+			panic("poly1305: unexpected overflow")
+		}
+
+		m0 := h0r0
+		m1 := add128(h1r0, h0r1) // These two additions don't overflow thanks again
+		m2 := add128(h2r0, h1r1) // to the 4 masked bits at the top of r0 and r1.
+		m3 := h2r1
+
+		t0 := m0.lo
+		t1, c := bitsAdd64(m1.lo, m0.hi, 0)
+		t2, c := bitsAdd64(m2.lo, m1.hi, c)
+		t3, _ := bitsAdd64(m3.lo, m2.hi, c)
+
+		// Now we have the result as 4 64-bit limbs, and we need to reduce it
+		// modulo 2¹³⁰ - 5. The special shape of this Crandall prime lets us do
+		// a cheap partial reduction according to the reduction identity
+		//
+		//     c * 2¹³⁰ + n  =  c * 5 + n  mod  2¹³⁰ - 5
+		//
+		// because 2¹³⁰ = 5 mod 2¹³⁰ - 5. Partial reduction since the result is
+		// likely to be larger than 2¹³⁰ - 5, but still small enough to fit the
+		// assumptions we make about h in the rest of the code.
+		//
+		// See also https://speakerdeck.com/gtank/engineering-prime-numbers?slide=23
+
+		// We split the final result at the 2¹³⁰ mark into h and cc, the carry.
+		// Note that the carry bits are effectively shifted left by 2, in other
+		// words, cc = c * 4 for the c in the reduction identity.
+		h0, h1, h2 = t0, t1, t2&maskLow2Bits
+		cc := uint128{t2 & maskNotLow2Bits, t3}
+
+		// To add c * 5 to h, we first add cc = c * 4, and then add (cc >> 2) = c.
+
+		h0, c = bitsAdd64(h0, cc.lo, 0)
+		h1, c = bitsAdd64(h1, cc.hi, c)
+		h2 += c
+
+		cc = shiftRightBy2(cc)
+
+		h0, c = bitsAdd64(h0, cc.lo, 0)
+		h1, c = bitsAdd64(h1, cc.hi, c)
+		h2 += c
+
+		// h2 is at most 3 + 1 + 1 = 5, making the whole of h at most
+		//
+		//     5 * 2¹²⁸ + (2¹²⁸ - 1) = 6 * 2¹²⁸ - 1
+	}
+
+	state.h[0], state.h[1], state.h[2] = h0, h1, h2
+}
+
+const (
+	maskLow2Bits    uint64 = 0x0000000000000003
+	maskNotLow2Bits uint64 = ^maskLow2Bits
+)
+
+// select64 returns x if v == 1 and y if v == 0, in constant time.
+func select64(v, x, y uint64) uint64 { return ^(v-1)&x | (v-1)&y }
+
+// [p0, p1, p2] is 2¹³⁰ - 5 in little endian order.
+const (
+	p0 = 0xFFFFFFFFFFFFFFFB
+	p1 = 0xFFFFFFFFFFFFFFFF
+	p2 = 0x0000000000000003
+)
+
+// finalize completes the modular reduction of h and computes
+//
+//     out = h + s  mod  2¹²⁸
+//
+func finalize(out *[TagSize]byte, h *[3]uint64, s *[2]uint64) {
+	h0, h1, h2 := h[0], h[1], h[2]
+
+	// After the partial reduction in updateGeneric, h might be more than
+	// 2¹³⁰ - 5, but will be less than 2 * (2¹³⁰ - 5). To complete the reduction
+	// in constant time, we compute t = h - (2¹³⁰ - 5), and select h as the
+	// result if the subtraction underflows, and t otherwise.
+
+	hMinusP0, b := bitsSub64(h0, p0, 0)
+	hMinusP1, b := bitsSub64(h1, p1, b)
+	_, b = bitsSub64(h2, p2, b)
+
+	// h = h if h < p else h - p
+	h0 = select64(b, h0, hMinusP0)
+	h1 = select64(b, h1, hMinusP1)
+
+	// Finally, we compute the last Poly1305 step
+	//
+	//     tag = h + s  mod  2¹²⁸
+	//
+	// by just doing a wide addition with the 128 low bits of h and discarding
+	// the overflow.
+	h0, c := bitsAdd64(h0, s[0], 0)
+	h1, _ = bitsAdd64(h1, s[1], c)
+
+	binary.LittleEndian.PutUint64(out[0:8], h0)
+	binary.LittleEndian.PutUint64(out[8:16], h1)
 }
diff --git a/poly1305/sum_noasm.go b/poly1305/sum_noasm.go
index 8a9c207..1682eda 100644
--- a/poly1305/sum_noasm.go
+++ b/poly1305/sum_noasm.go
@@ -6,10 +6,7 @@
 
 package poly1305
 
-// Sum generates an authenticator for msg using a one-time key and puts the
-// 16-byte result into out. Authenticating two different messages with the same
-// key allows an attacker to forge messages at will.
-func Sum(out *[TagSize]byte, msg []byte, key *[32]byte) {
+func sum(out *[TagSize]byte, msg []byte, key *[32]byte) {
 	h := newMAC(key)
 	h.Write(msg)
 	h.Sum(out)
diff --git a/poly1305/sum_ppc64le.go b/poly1305/sum_ppc64le.go
index 2402b63..3233616 100644
--- a/poly1305/sum_ppc64le.go
+++ b/poly1305/sum_ppc64le.go
@@ -7,62 +7,52 @@
 package poly1305
 
 //go:noescape
-func initialize(state *[7]uint64, key *[32]byte)
+func update(state *macState, msg []byte)
 
-//go:noescape
-func update(state *[7]uint64, msg []byte)
-
-//go:noescape
-func finalize(tag *[TagSize]byte, state *[7]uint64)
-
-// Sum generates an authenticator for m using a one-time key and puts the
-// 16-byte result into out. Authenticating two different messages with the same
-// key allows an attacker to forge messages at will.
-func Sum(out *[16]byte, m []byte, key *[32]byte) {
+func sum(out *[16]byte, m []byte, key *[32]byte) {
 	h := newMAC(key)
 	h.Write(m)
 	h.Sum(out)
 }
 
 func newMAC(key *[32]byte) (h mac) {
-	initialize(&h.state, key)
+	initialize(key, &h.r, &h.s)
 	return
 }
 
-type mac struct {
-	state [7]uint64 // := uint64{ h0, h1, h2, r0, r1, pad0, pad1 }
+// mac is a wrapper for macGeneric that redirects calls that would have gone to
+// updateGeneric to update.
+//
+// Its Write and Sum methods are otherwise identical to the macGeneric ones, but
+// using function pointers would carry a major performance cost.
+type mac struct{ macGeneric }
 
-	buffer [TagSize]byte
-	offset int
-}
-
-func (h *mac) Write(p []byte) (n int, err error) {
-	n = len(p)
+func (h *mac) Write(p []byte) (int, error) {
+	nn := len(p)
 	if h.offset > 0 {
-		remaining := TagSize - h.offset
-		if n < remaining {
-			h.offset += copy(h.buffer[h.offset:], p)
-			return n, nil
+		n := copy(h.buffer[h.offset:], p)
+		if h.offset+n < TagSize {
+			h.offset += n
+			return nn, nil
 		}
-		copy(h.buffer[h.offset:], p[:remaining])
-		p = p[remaining:]
+		p = p[n:]
 		h.offset = 0
-		update(&h.state, h.buffer[:])
+		update(&h.macState, h.buffer[:])
 	}
-	if nn := len(p) - (len(p) % TagSize); nn > 0 {
-		update(&h.state, p[:nn])
-		p = p[nn:]
+	if n := len(p) - (len(p) % TagSize); n > 0 {
+		update(&h.macState, p[:n])
+		p = p[n:]
 	}
 	if len(p) > 0 {
 		h.offset += copy(h.buffer[h.offset:], p)
 	}
-	return n, nil
+	return nn, nil
 }
 
 func (h *mac) Sum(out *[16]byte) {
-	state := h.state
+	state := h.macState
 	if h.offset > 0 {
 		update(&state, h.buffer[:h.offset])
 	}
-	finalize(out, &state)
+	finalize(out, &state.h, &state.s)
 }
diff --git a/poly1305/sum_ppc64le.s b/poly1305/sum_ppc64le.s
index 55c7167..4e20bf2 100644
--- a/poly1305/sum_ppc64le.s
+++ b/poly1305/sum_ppc64le.s
@@ -58,7 +58,6 @@
 GLOBL ·poly1305Mask<>(SB), RODATA, $16
 
 // func update(state *[7]uint64, msg []byte)
-
 TEXT ·update(SB), $0-32
 	MOVD state+0(FP), R3
 	MOVD msg_base+8(FP), R4
@@ -180,68 +179,3 @@
 	MOVD R9, 8(R3)
 	MOVD R10, 16(R3)
 	RET
-
-// func initialize(state *[7]uint64, key *[32]byte)
-TEXT ·initialize(SB), $0-16
-	MOVD state+0(FP), R3
-	MOVD key+8(FP), R4
-
-	// state[0...7] is initialized with zero
-	// Load key
-	MOVD 0(R4), R5
-	MOVD 8(R4), R6
-	MOVD 16(R4), R7
-	MOVD 24(R4), R8
-
-	// Address of key mask
-	MOVD $·poly1305Mask<>(SB), R9
-
-	// Save original key in state
-	MOVD R7, 40(R3)
-	MOVD R8, 48(R3)
-
-	// Get mask
-	MOVD (R9), R7
-	MOVD 8(R9), R8
-
-	// And with key
-	AND R5, R7, R5
-	AND R6, R8, R6
-
-	// Save masked key in state
-	MOVD R5, 24(R3)
-	MOVD R6, 32(R3)
-	RET
-
-// func finalize(tag *[TagSize]byte, state *[7]uint64)
-TEXT ·finalize(SB), $0-16
-	MOVD tag+0(FP), R3
-	MOVD state+8(FP), R4
-
-	// Get h0, h1, h2 from state
-	MOVD 0(R4), R5
-	MOVD 8(R4), R6
-	MOVD 16(R4), R7
-
-	// Save h0, h1
-	MOVD  R5, R8
-	MOVD  R6, R9
-	MOVD  $3, R20
-	MOVD  $-1, R21
-	SUBC  $-5, R5
-	SUBE  R21, R6
-	SUBE  R20, R7
-	MOVD  $0, R21
-	SUBZE R21
-
-	// Check for carry
-	CMP  $0, R21
-	ISEL $2, R5, R8, R5
-	ISEL $2, R6, R9, R6
-	MOVD 40(R4), R8
-	MOVD 48(R4), R9
-	ADDC R8, R5
-	ADDE R9, R6
-	MOVD R5, 0(R3)
-	MOVD R6, 8(R3)
-	RET
diff --git a/poly1305/sum_s390x.go b/poly1305/sum_s390x.go
index ec99e07..a8920ee 100644
--- a/poly1305/sum_s390x.go
+++ b/poly1305/sum_s390x.go
@@ -22,10 +22,7 @@
 //go:noescape
 func poly1305vmsl(out *[16]byte, m *byte, mlen uint64, key *[32]byte)
 
-// Sum generates an authenticator for m using a one-time key and puts the
-// 16-byte result into out. Authenticating two different messages with the same
-// key allows an attacker to forge messages at will.
-func Sum(out *[16]byte, m []byte, key *[32]byte) {
+func sum(out *[16]byte, m []byte, key *[32]byte) {
 	if cpu.S390X.HasVX {
 		var mPtr *byte
 		if len(m) > 0 {