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

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


在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

posted on 2022-07-27 12:02  ahfuzhang  阅读(152)  评论(0)    收藏  举报