【小优化】golang中取两个字符串的公共前缀的长度

作者:张富春(ahfuzhang)，转载时请注明作者和引用链接，谢谢！

• cnblogs博客
• zhihu
• Github
• 公众号:一本正经的瞎扯
  
  

---

在VM的merge部分的代码中发现这样一个函数：

func commonPrefixLen(a, b []byte) int {
    i := 0
    if len(a) > len(b) {
        for i < len(b) && a[i] == b[i] {
            i++
        }
    } else {
        for i < len(a) && a[i] == b[i] {
            i++
        }
    }
    return i
}

在merge整个sstable的时候，需要算出所有字符串的公共前缀，然后存储的时候就不再存储公共前缀的部分，这样就实现了进一步的数据压缩。

这个函数调用得非常频繁，于是想尝试用SIMD来优化。

后来试了一下，简单的写法就能提升1.6倍，优化的代码如下：

// macbook pro 2019, Intel(R) Core(TM) i7-6700HQ CPU @ 2.60GHz
// 15.70 ns/op
func commonPrefixLenOneByOne(a, b []byte) int {
    i := 0
    if len(a) > len(b) {
        for i < len(b) && a[i] == b[i] {
            i++
        }
    } else {
        for i < len(a) && a[i] == b[i] {
            i++
        }
    }
    return i
}

// from this issue: https://github.com/VictoriaMetrics/VictoriaMetrics/issues/2254
// macbook pro 2019, Intel(R) Core(TM) i7-6700HQ CPU @ 2.60GHz
// 9.785 ns/op when memory aligned, 1.6 times faster
// 9.833 ns/op when memory not aligned
func commonPrefixLen(a, b []byte) int {
    if len(a) > len(b) {
        a, b = b, a
    }
    if len(a) < 8 {
        return commonPrefixLenOneByOne(a, b)
    }
    const size = 8
    compareTimes := len(a) / size
    addrA := uintptr(unsafe.Pointer(&a[0]))
    addrB := uintptr(unsafe.Pointer(&b[0]))
    for i := 0; i < compareTimes; i++ {
        v1 := (*uint64)(unsafe.Pointer(addrA))
        v2 := (*uint64)(unsafe.Pointer(addrB))
        if *v1 != *v2 {
            return i*size + commonPrefixLenOneByOne(a[i*size:], b[i*size:])
        }
        addrA += size
        addrB += size
    }
    return compareTimes*size + commonPrefixLenOneByOne(a[compareTimes*size:], b[compareTimes*size:])
}

思路就是用unsafe指针把连续的8个byte转换为uint64的指针，这样就可以一条指令比较8个字节。

这个优化已经提了一个PR到VM的仓库：https://github.com/VictoriaMetrics/VictoriaMetrics/pull/2913