-
Notifications
You must be signed in to change notification settings - Fork 422
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
sstable: generate streamlined blockIter code
We automatically generate a "streamlined" version of the blockIter code (in `block_iter_streamlined.gen.go`). The streamlined code can be used in the common case and does not support all features. This improves performance by reducing the conditionals in the hot path. The streamlinedBlockIter constant can be used in if statements; in the streamlined version, it gets replaced by "true" (allowing the compiler to omit blocks of code from the streamlined version).
- Loading branch information
1 parent
98a8ea9
commit ee9934e
Showing
5 changed files
with
1,134 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,199 @@ | ||
| // Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use | ||
| // of this source code is governed by a BSD-style license that can be found in | ||
| // the LICENSE file. | ||
|
|
||
| // This program is used to generate a "streamlined" version of the code in | ||
| // block_iter.go. The goal is to remove conditionals from the hot path in the | ||
| // common case. | ||
| // | ||
| // The block_iter.go code uses a special constant streamlinedBlockIter which is | ||
| // always false in th code. This program generates streamlined versions of the | ||
| // original methods in which the constant is replaced with true. The methods | ||
| // that need streamlining are those that use the constant or call (directly or | ||
| // indirectly) a method which uses it. | ||
| // | ||
| // The higher level code decides whether to call the general or streamlined | ||
| // version of a method. | ||
| package main | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "cmp" | ||
| "go/ast" | ||
| "go/parser" | ||
| "go/printer" | ||
| "go/token" | ||
| "os" | ||
| "slices" | ||
| "strings" | ||
| ) | ||
|
|
||
| const inputFile = "block_iter.go" | ||
| const outputFile = "block_iter_streamlined.gen.go" | ||
| const specialConstant = "streamlinedBlockIter" | ||
| const excludedMethod = "init" | ||
|
|
||
| const header = `// Code generated by block-iter-codegen; DO NOT EDIT. | ||
| package sstable | ||
| import ( | ||
| "bytes" | ||
| "encoding/binary" | ||
| "slices" | ||
| "unsafe" | ||
| "github.com/cockroachdb/errors" | ||
| "github.com/cockroachdb/pebble/internal/base" | ||
| "github.com/cockroachdb/pebble/internal/invariants" | ||
| "github.com/cockroachdb/pebble/internal/manual" | ||
| ) | ||
| ` | ||
|
|
||
| type methodInfo struct { | ||
| node *ast.FuncDecl | ||
| callers []string | ||
| needsStreamlining bool | ||
| } | ||
|
|
||
| // methods is keyed on the method name. | ||
| var methods = map[string]*methodInfo{} | ||
|
|
||
| func main() { | ||
| // Read the source file. | ||
| fset := token.NewFileSet() | ||
| astFile, err := parser.ParseFile(fset, inputFile, nil, parser.ParseComments|parser.SkipObjectResolution) | ||
| if err != nil { | ||
| panic(err) | ||
| } | ||
|
|
||
| CollectMethods(astFile) | ||
|
|
||
| delete(methods, excludedMethod) | ||
|
|
||
| // Populate the caller-callee relationships. Conceptually we build a graph | ||
| // where each relationship is represented as an edge from callee to caller. | ||
| for _, m := range methods { | ||
| RecordCaller(m.node) | ||
| } | ||
|
|
||
| // Mark all methods that directly or indirectly use the special constant. | ||
| for name, m := range methods { | ||
| if UsesSpecialConstant(m.node) { | ||
| DF(name) | ||
| } | ||
| } | ||
|
|
||
| var toOutput []*methodInfo | ||
| for name, m := range methods { | ||
| if m.needsStreamlining { | ||
| // Rename method. | ||
| m.node.Name.Name = Streamlined(name) | ||
|
|
||
| // Rename uses of the special constant and any calls to methods that need | ||
| // streamlining. | ||
| ast.Inspect(astFile, func(n ast.Node) bool { | ||
| switch n := n.(type) { | ||
| case *ast.Ident: | ||
| if n.Name == specialConstant { | ||
| // This is kind of hacky, as "true" wouldn't be an Ident; but it prints | ||
| // out correctly. | ||
| n.Name = "true" | ||
| } | ||
| case *ast.CallExpr: | ||
| // Rename method calls to streamlined functions. | ||
| if s, ok := n.Fun.(*ast.SelectorExpr); ok { | ||
| if m, ok := methods[s.Sel.Name]; ok && m.needsStreamlining { | ||
| s.Sel.Name = Streamlined(s.Sel.Name) | ||
| } | ||
| } | ||
| } | ||
| return true | ||
| }) | ||
|
|
||
| toOutput = append(toOutput, m) | ||
| } | ||
| } | ||
| // Sort by position in the original file. | ||
| slices.SortFunc(toOutput, func(a, b *methodInfo) int { | ||
| return cmp.Compare(a.node.Body.Lbrace, b.node.Body.Lbrace) | ||
| }) | ||
|
|
||
| var buf bytes.Buffer | ||
| buf.WriteString(header) | ||
| for _, m := range toOutput { | ||
| buf.WriteString("\n") | ||
| printer.Fprint(&buf, fset, m.node) | ||
| buf.WriteString("\n") | ||
| } | ||
| if err := os.WriteFile(outputFile, buf.Bytes(), 0666); err != nil { | ||
| panic(err) | ||
| } | ||
| } | ||
|
|
||
| // CollectMethods adds all top-level method declarations to the methods map. | ||
| func CollectMethods(root ast.Node) { | ||
| ast.Inspect(root, func(n ast.Node) bool { | ||
| switch n := n.(type) { | ||
| case *ast.FuncDecl: | ||
| if n.Recv != nil { | ||
| methods[n.Name.Name] = &methodInfo{node: n} | ||
| } | ||
| return false | ||
| } | ||
| return true | ||
| }) | ||
| } | ||
|
|
||
| // RecordCaller adds the given function to the callers lists of any methods it | ||
| // calls directly. | ||
| func RecordCaller(f *ast.FuncDecl) { | ||
| fName := f.Name.Name | ||
| ast.Inspect(f, func(n ast.Node) bool { | ||
| switch n := n.(type) { | ||
| case *ast.CallExpr: | ||
| if s, ok := n.Fun.(*ast.SelectorExpr); ok { | ||
| if m, ok := methods[s.Sel.Name]; ok { | ||
| m.callers = append(m.callers, fName) | ||
| } | ||
| } | ||
| } | ||
| return true | ||
| }) | ||
| } | ||
|
|
||
| // UsesSpecialConstant returns true if the given function uses the special | ||
| // constant directly. | ||
| func UsesSpecialConstant(f *ast.FuncDecl) bool { | ||
| result := false | ||
| ast.Inspect(f, func(n ast.Node) bool { | ||
| switch n := n.(type) { | ||
| case *ast.Ident: | ||
| if n.Name == specialConstant { | ||
| result = true | ||
| } | ||
| } | ||
| return true | ||
| }) | ||
| return result | ||
| } | ||
|
|
||
| // DF recursively sets the needsStreamlining flag for the given method and all | ||
| // direct or indirect caller methods. | ||
| func DF(name string) { | ||
| if m, ok := methods[name]; ok && !m.needsStreamlining { | ||
| m.needsStreamlining = true | ||
| for _, caller := range m.callers { | ||
| DF(caller) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Streamlined returns the streamlined version of the method name. | ||
| // E.g. "foo" -> "streamlinedFoo"; "Foo" -> "StreamlinedFoo". | ||
| func Streamlined(n string) string { | ||
| if n[0] >= 'A' && n[0] <= 'Z' { | ||
| return "Streamlined" + n | ||
| } | ||
| return "streamlined" + strings.ToUpper(n[:1]) + n[1:] | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.