Golang 中字典的 Comma Ok 是如何實現的

h3l-github發表於2020-03-01

眾所周知，Golang 中函式的返回值的數量是固定的，而不是像 Python 中那樣，函式的返回值數量是不固定的。

如果我們把 Golang 中對 map 的取值看作是一個函式的話，那麼直接取值和用 comma ok 方式取值的實現就變得很意思。

Golang 中 map 的取值方式

v1, ok := m["test"]
v2 := m2["test"]

先看看彙編是如何實現的。

package main

import "log"

func main() {
    m1 := make(map[string]string)
    v1, ok := m1["test"]
    v2 := m1["test"]

    log.Println(v1, v2, ok)
}

儲存上述檔案為 map_test.go，執行go tool compile -S map_test.go，擷取關鍵部分

...
    0x00a9 00169 (map_test.go:7)    CALL    runtime.mapaccess2_faststr(SB)
...
    0x00f8 00248 (map_test.go:8)    CALL    runtime.mapaccess1_faststr(SB)
...

可以看到，雖然都是 m1["test"]，但是卻呼叫了 runtime 中不同的方法。可以在 go/src/runtime/map_faststr.go 檔案中看到

func mapaccess2_faststr(t *maptype, h *hmap, ky string) (unsafe.Pointer, bool) {}
func mapaccess1_faststr(t *maptype, h *hmap, ky string) unsafe.Pointer {}

這樣明顯就對上了，但是 Golang 又是如何實現把 m["test"] 替換為 mapaccess2_faststr 或者 mapaccess1_faststr 的呢？

這就涉及 Golang 的編譯過程了。檢視官方文件，我們知道編譯的過程包括：

Parsing，包括詞法分析，語法分析，抽象語法樹的生成。
Type-checking and AST transformations，包括型別檢查，抽象語法樹轉換。
Generic SSA，中間程式碼生成
Generating machine code，生成機器碼

現在我們就一步一步的看一看，m["test"]是如何變成mapaccess2_faststr的。（mapaccess1_faststr同理，故不贅述）

詞法分析

詞法分析，Golang 中的詞法分析主要是通過go/src/cmd/compile/internal/syntax/scanner.go（簡稱 scanner.go）與 go/src/cmd/compile/internal/syntax/tokens.go（簡稱 tokens.go）完成的，其中，tokens.go 中定義各種字元會被轉化成什麼樣。例如： tokens.go 中分別定義了 [ 與 ]

_Lbrack    // [
_Rbrack    // ]

會被怎樣處理。

而在 scanner.go 中，通過一個大的 switch 處理各種字元。處理 [ 與 ] 的部分程式碼如下:

switch c {
    // 略過
    case '[':
        s.tok = _Lbrack
    case ']':
        s.nlsemi = true
        s.tok = _Rbrack
    // 略過
}

語法分析

語法分析階段會將詞法分析階段生成的轉換成各種 Expr（表示式），表示式的定義在go/src/cmd/compile/internal/syntax/nodes.go（簡稱 nodes.go）。而 map 取值的表示式定義如下：

// X[Index]
IndexExpr struct {
    X     Expr
    Index Expr
    expr
}

之後再通過go/src/cmd/compile/internal/syntax/parser.go（簡稱 parser.go）中的 pexpr 函式將詞法分析階段的 token 轉化為表示式。關鍵部分如下：

switch p.tok {
    // 略
    case _Lbrack: // 遇到一個左方括號
        p.next()
        p.xnest++

        var i Expr
        if p.tok != _Colon { // 遇到一個右方括號
            i = p.expr()
            if p.got(_Rbrack) {
                // x[i]
                t := new(IndexExpr) // 生成一個 Index表示式
                t.pos = pos
                t.X = x
                t.Index = i
                x = t
                p.xnest--
                break
            }
        }
    //略
}

至此，已經將 m["key"] 轉化為一個 IndexExpr 了。

抽象語法樹生成

之後，在go/src/cmd/compile/internal/gc/noder.go檔案中，再將 IndexExpr 轉化成一個OINDEX型別的 node，關鍵程式碼如下：

switch expr := expr.(type) {
    // 略
    case *syntax.IndexExpr:
        return p.nod(expr, OINDEX, p.expr(expr.X), p.expr(expr.Index))
    // 略
}

其中各種操作型別的定義，如上述的OINDEX在檔案go/src/cmd/compile/internal/gc/syntax.go(簡稱為 syntax.go) 中，如下

OINDEX       // Left[Right] (index of array or slice)

型別檢查

對於上文獲得的最後一個 OINDEX 型別的 node，他取值的物件即可能是字典，也可能是陣列、字串等。所以要對他們進行區分，而型別檢查部分就是做這方面工作的。跟本文相關的函式是go/src/cmd/compile/internal/gc/typecheck.go（簡稱為 typecheck.go）檔案中的typecheck1函式。其中關鍵程式碼如下：

func typecheck1(n *Node, top int) (res *Node) {
    // 略
    switch n.Op {
    case OINDEX: // 處理 OINDEX 型別的節點
        // 略過部分檢查程式碼
        // 獲取 Left[Right] 中的 Left的型別
        l := n.Left
        t := l.Type
        switch t.Etype {
        default:
            yyerror("invalid operation: %v (type %v does not support indexing)", n, t)
            n.Type = nil
            return n

        case TSTRING, TARRAY, TSLICE:
            // 處理 Left 是字串、陣列、切片的情況
            // 略

        case TMAP:
            // 如果 Left 是 MAP，則把該 node 的操作變成 OINDEXMAP
            n.Right = defaultlit(n.Right, t.Key())
            if n.Right.Type != nil {
                n.Right = assignconv(n.Right, t.Key(), "map index")
            }
            n.Type = t.Elem()
            n.Op = OINDEXMAP
            n.ResetAux()
        }
    }
}

繼續對操作為OINDEXMAP（OINDEXMAP也定義在syntax.go中）的 node 節點進行分析。可以看到，在typecheck.go的typecheckas2函式中，繼續對OINDEXMAP的節點進行分析。其中關鍵程式碼如下：

func typecheckas2(n *Node) {
    // 略
    cl := n.List.Len()
    cr := n.Rlist.Len()
    // 略

    // x, ok = y
    // 引數左邊是兩個，右邊是一個
    if cl == 2 && cr == 1 {
        switch r.Op {
        case OINDEXMAP, ORECV, ODOTTYPE:
            switch r.Op {
            case OINDEXMAP:
                // 如果操作的物件是OINDEXMAP，將其變為 OAS2MAPR
                n.Op = OAS2MAPR
            }
        }
    }
    //略
}

最終，我們的v1, ok := m["test"]的語句，變成了一個型別為OAS2MAPR的語法樹節點。

中間程式碼生成

中間程式碼生成即將語法樹生成與機器碼無關的中間程式碼。生成中間程式碼的檔案為go/src/cmd/compile/internal/gc/walk.go(簡稱 walk.go)，與本文相關的為walk.go檔案中的walkexpr函式。關鍵程式碼如下：

func walkexpr(n *Node, init *Nodes) *Node {
    switch n.Op {
    // a,b = m[i]
    case OAS2MAPR:
        // 略

        // from:
        //   a,b = m[i]
        // to:
        //   var,b = mapaccess2*(t, m, i)
        //   a = *var
        a := n.List.First()

        // 根據 map 中 key 值型別不同以及值的長度進行優化
        if w := t.Elem().Width; w <= 1024 { // 1024 must match runtime/map.go:maxZero
            fn := mapfn(mapaccess2[fast], t)
            r = mkcall1(fn, fn.Type.Results(), init, typename(t), r.Left, key)
        } else {
            fn := mapfn("mapaccess2_fat", t)
            z := zeroaddr(w)
            r = mkcall1(fn, fn.Type.Results(), init, typename(t), r.Left, key, z)
        }
        // 略
        n.Rlist.Set1(r)
        n.Op = OAS2FUNC

       // 略

        n = typecheck(n, ctxStmt)
        n = walkexpr(n, init)
    }
}