【go源码分析】strings.go 里的那些骚操作

go version go1.11 darwin/amd64
/src/strings/strings.go

strings.go 文件中定义了近40个常用的字符串操作函数（公共函数）。以下是主要的几个函数。

以及 ToUpper ToLower Title（Title 函数把单词转换成标题形式，不是ToTitle）。

还有一些以上函数派生出的其他函数。比如：Contains 基本是通过 Index 函数实现的；与 Index 原理一致的 LastIndex 函数；与 Trim 有关的 TrimLeft TrimRight 等。

接下来，本文会对 Index Trim Join Repeat Replace 函数进行分析。

ps: len 返回的是字符串的字节数，不是字符数。字符数请使用 utf8.RuneCountInString

Index: RabinKarp 算法实现

Index(s, substr string) int, 返回 substr 在 s 中第一次出现的位置，不存在返回 -1；采用RabinKarp算法

Index 函数会先对 substr 的长度 n 进行判断，对特殊情况做快速处理。

其次，如果长度 n 以及 len(s) 足够小，则使用BruteForce算法：即暴力匹配，拿 substr 与 s[i:i+n] 进行比较，如果相等，返回 i，其中 i = (from 0 to len(s) - n)...

最后，会先尝试暴力匹配，如果匹配失败次数超过临界点，则换成 RabinKarp 算法。

（为了方便阅读，文中不放全部代码，只展示核心代码与部分结构代码）

Index

func Index(s, substr string) int {
    n := len(substr)  // len 返回的是字节数
    switch {
    case n == 0:
        return 0
    case n == 1:
        // substr 是单字节字符串，则直接单字节进行比较
        return IndexByte(s, substr[0])
    case n == len(s):
        if substr == s {
            return 0
        }
        return -1
    case n > len(s):
        return -1
    case n <= bytealg.MaxLen:
        // Use brute force when s and substr both are small
        if len(s) <= bytealg.MaxBruteForce {
            return bytealg.IndexString(s, substr)
        }
        
        // 这里有一大段省略的代码
        // 循环尝试 substr == s[i:i+n]
        // 如果失败次数过多，则使用 bytealg.IndexString(s, substr)
    }
    
    // 这里有一大段省略的代码
    // 循环尝试 substr == s[i:i+n]
    // 如果失败次数过多，则使用 indexRabinKarp(s[i:], substr)
    
    t := s[:len(s)-n+1]
    for i < len(t) {
        
        // ... 省略代码
        
        // 如果失败次数过多，则使用 RabinKarp 算法
        if fails >= 4+i>>4 && i < len(t) {
            // 这里使用 s[i:] 作为参数
            // 是因为前面的 s[:i] 都已经尝试过了
            j := indexRabinKarp(s[i:], substr)
            if j < 0 {
                return -1
            }
            return i + j
        }
    }
    
    return -1
}

在看 indexRabinKarp 函数之前，我们先了解一下 RabinKarp 算法。

RobinKarp算法是由 Robin 和 Karp 提出的字符串匹配算法。该算法在实际应用中有较好的表现。

算法的核心步骤：

• const primeRK = 16777619 // 大素数
• 对 substr 构造 hash 值。 n = len(substr)，hash = (substr[0]*pow(primeRK, n-1) + substr[1]*pow(primeRK, n-2) + ... + substr[n-1]*pow(primeRK, 0)) % anotherBiggerPrime
• 对 s 的每 n 个子串按照相同逻辑构造 hash 值，判断与 substr 的 hash 是否相等；如果 hash 相等，则比较子串是否真的与 substr 相等
• 重复第三步，直到找到，或者未找到。

ps:

• 该算法之所以快，是因为 s[i+1, i+n+1] 的 hash 值可以由 s[i, i+n] 的 hash 值计算出。即h(i+1) = ((h(i) - s[i] * pow(primeRK, n-1)) * primeRK + s[i+n+1]) % anotherBiggerPrime
• 另外，go 计算 hash 时并没有 % anotherBiggerPrime，而是定义了 hash 为 uint32 类型，利用整型溢出实现了对 2**32 取模的效果。（一般来说是对另一个大素数取模，显然这里不是，不过毕竟这么大的数也没啥大影响）

该算法预处理时间为 O(n)，n 为 len(substr)，运行最坏时间为 O((n-m+1)m)，m 为 len(s)。最坏情况为每个子串的 hash 都与 substr 的一样。在平均情况下，运行时间还是很好的。

除了 RabinKarp 算法外，还要一些其他的字符串匹配算法。《算法》导论中介绍了另外两种优秀的算法，分别是 有限自动机 与 Knuth-Morris-Pratt 算法（即 KMP 算法），这两个算法的运行时间都为 O(m)。

下面是 indexRabinKarp 函数

indexRabinKarp

func indexRabinKarp(s, substr string) int {
    // Rabin-Karp search
    // hashss 是 substr 根据上述方法计算出的 hash 值
    // pow 是 primeRK 的 len(substr) 次幂
    hashss, pow := hashStr(substr)
    n := len(substr)

    // 计算 s[:n] 的 hash 值
    var h uint32
    for i := 0; i < n; i++ {
        h = h*primeRK + uint32(s[i])
    }
    if h == hashss && s[:n] == substr {
        return 0
    }

    for i := n; i < len(s); {
        // 计算下一个子串的 hash 值
        h *= primeRK
        h += uint32(s[i])
        h -= pow * uint32(s[i-n])
        i++
        
        // 如果 hash 相等 且子串相等，则返回对应下标
        if h == hashss && s[i-n:i] == substr {
            return i - n
        }
    }
    return -1
}

hashStr 函数跟计算 s[:n] 的 逻辑一致。不过不得不提一下 pow 的计算方法。

hashStr

func hashStr(sep string) (uint32, uint32) {
    hash := uint32(0)
    for i := 0; i < len(sep); i++ {
        hash = hash*primeRK + uint32(sep[i])
    }
    
    // 下面我用 python 的乘方元素符 ** 代表乘方
    // 我们已 len(sep) 为 6 为例来看此函数
    // 6 的二进制是 110
    // 每次循环，pow 和 sq 分别为
    // i: 110  pow: 1  sq: rk
    // i: 11   pow: 1  sq: rk ** 2
    // i: 1    pow: 1 * (rk ** 2)  sq: rk ** 4
    // i: 0    pow: 1* (rk ** 2) * (rk ** 4)  sq: rk ** 8
    // pow: 1* (rk ** 2) * (rk ** 4) = 1 * (rk ** 6) 即是 pow(rk, 6)
    var pow, sq uint32 = 1, primeRK
    for i := len(sep); i > 0; i >>= 1 {
        if i&1 != 0 {
            pow *= sq
        }
        sq *= sq
    }
    return hash, pow
}

以上是 Index 函数的实现逻辑。

Trim: 出神入化的位操作

Trim(s string, cutset string) string 返回去除首尾存在于 cutset 的字符的切片。

执行 fmt.Println(strings.Trim("hello world", "hld"))
输出 ello wor

Trim的本质逻辑也比较简单：

• 根据 cutset 构造一个函数，该函数签名为 func(rune) bool，接受一个 rune类型的值，返回该值是否在 cutset 中
• 然后调用 TrimLeft TrimRight；这两个函数调用了 indexFunc，其逻辑也比较简单，不再赘述

函数 makeCutsetFunc(cutset string) func(rune) bool 就是刚才提到的构造 判断 rune 值是否在 cutset 中的函数 的函数。

makeCutsetFunc

func makeCutsetFunc(cutset string) func(rune) bool {
    // 如果 cutset 是单个字符，则直接返回一个简单函数，
    // 该函数判断入参 r 是否与 cutset[0] 相等
    if len(cutset) == 1 && cutset[0] < utf8.RuneSelf {
        return func(r rune) bool {
            return r == rune(cutset[0])
        }
    }

    // 如果 cutset 全是 ASCII 码
    // 则使用构造的 as (asciiSet类型)判断 rune 是否在 cutset 中
    if as, isASCII := makeASCIISet(cutset); isASCII {
        return func(r rune) bool {
            return r < utf8.RuneSelf && as.contains(byte(r))
        }
    }

    // 调用 IndexRune 方法判断 r 是否在 cutset 中
    // IndexRune 其实就是 Index 的变种
    return func(r rune) bool { return IndexRune(cutset, r) >= 0 }
}

其中，最有意思的要数 makeASCIISet 函数，该函数用了一个 [8]uint32 数组实现了 128 个 ASCII 码的 hash 表。

asciiSet

// asciiSet 一共 32 个字节，一共 256 位，
// 其中低 128 位分别代表了 128 个 ascii 码
type asciiSet [8]uint32

// makeASCIISet creates a set of ASCII characters and reports whether all
// characters in chars are ASCII.
func makeASCIISet(chars string) (as asciiSet, ok bool) {
    for i := 0; i < len(chars); i++ {
        c := chars[i]
        // const utf8.RuneSelf = 0x80
        // 小于 utf8.RuneSelf 的值是 ASCII 码
        // 大于 utf8.RuneSelf 的值是其他 utf8 字符的部分
        if c >= utf8.RuneSelf {
            return as, false
        }
        
        // ASCII 的范围是 0000 0000 - 0111 1111
        // c >> 5 的范围是 000 - 011，即最大为 3
        // 31 的二进制是 0001 1111
        // 1 << uint(c&31) 的结果刚好也在 uint 范围内
        
        as[c>>5] |= 1 << uint(c&31)
    }
    return as, true
}

// contains reports whether c is inside the set.
// 为了兼容入参 c 为 byte 类型， c >> 5 < 8
// 所以 asciiSet 类型为 [8]uint32，数组长度为 8
// 否则如果只考虑 128 个 ascii 码的话，[4]uint32 就够了
func (as *asciiSet) contains(c byte) bool {
    return (as[c>>5] & (1 << uint(c&31))) != 0
}

以上是 Trim 函数及其位操作。

从 Join Repeat Replace 看字符串如何生成

这三个函数的逻辑都很简单，不再赘述。

频繁申请内存是很耗费时间的，所以在生成某个字符串时，如果能够预知字符串的长度，就能直接申请对应长度的内存，然后调用 copy(dst, src []Type) int 函数把字符复制到对应位置，最后把 []byte 强转成字符串类型即可。

Join

func Join(a []string, sep string) string {
    // 省略了部分特殊情况处理的代码

    // 计算目标字符串总长度
    n := len(sep) * (len(a) - 1)
    for i := 0; i < len(a); i++ {
        n += len(a[i])
    }
    
    // 申请内存
    b := make([]byte, n)
    bp := copy(b, a[0])
    // 复制内容
    for _, s := range a[1:] {
        bp += copy(b[bp:], sep)
        bp += copy(b[bp:], s)
    }
    // 返回数据
    return string(b)
}

Repeat

func Repeat(s string, count int) string {
    // 特殊情况处理
    if count < 0 {
        panic("strings: negative Repeat count")
    } else if count > 0 && len(s)*count/count != len(s) {
        panic("strings: Repeat count causes overflow")
    }

    b := make([]byte, len(s)*count)
    bp := copy(b, s)

    // 2倍，4倍，8倍的扩大，直到 bp 不小于目标长度
    for bp < len(b) {
        copy(b[bp:], b[:bp])
        bp *= 2
    }
    return string(b)
}

Replace

func Replace(s, old, new string, n int) string {
    // 省略一下特殊情况代码

    // 计算新字符串的长度
    t := make([]byte, len(s)+n*(len(new)-len(old)))
    w := 0
    start := 0
    for i := 0; i < n; i++ {
        j := start
        if len(old) == 0 {
            if i > 0 {
                _, wid := utf8.DecodeRuneInString(s[start:])
                j += wid
            }
        } else {
            j += Index(s[start:], old)
        }
        // 查找旧字符串的位置，复制
        w += copy(t[w:], s[start:j])
        w += copy(t[w:], new)
        start = j + len(old)
    }
    w += copy(t[w:], s[start:])
    return string(t[0:w])
}

以上是关于生成字符串时避免多次分配内存的高效做法。