POC JIT with Go (Plugins)

POC JIT With Go (plugins) | xnacly - blog POC JIT With Go (plugins) Jan 5, 2024 834 Words 4 Minute read Tags: Go I recently implemented a very rudimentary JIT compiler in go. This compiler generates go source code, compiles it by invoking the go compiler and loading the resulting shared object into the current process and invoking the compiled function. The scope is currently limited to arithmetic expressions. I implemented a tree-walk interpreter as well as a byte code compiler and virtual machine to compare this JIT to. The front-end (lexer, parser) is shared, the back-end is specific to the approach of evaluation. As mentioned before the scope I choose for this compiler is limited to arithmetic expressions, thus I want to accept the following: TEXT 1 1+1.2-5 Tokenizing We convert our character stream to tokens: TEXT 1 NUMBER 1 2 PLUS 3 NUMBER 1.2 4 MINUS 5 NUMBER 5 Parsing We parse the tokens and produce an abstract syntax tree via recursive descent: TEXT 1 Binary { 2 Type: MINUS, 3 Left: Binary { 4 Type: PLUS, 5 Left: Number { 6 Value: 1 7 } 8 Right: Number { 9 Value: 1.2 10 } 11 } 12 Right: Number { 13 Value: 5 14 } 15 } Code-generation I want to generate go code from the abstract syntax tree, I know for arithmetics its pretty idiotic, but i want to evaluate the pipeline in a comparable way to byte code compilation+evaluation and tree-walk interpreting. Thus we generate the following go source code: GO 1 package main 2 func Main () float64 { return 1E+00 + 1.2E+00 - 5E+00 } Tip The generated function has to be exported, otherwise the plugin package will not recognize it. I had to choose this weird way of representing floating point integers, because otherwise there would be some weird results caused by go not casting numbers to floating point integers. Compilation The go compiler sits in an internal go package, see its repo here and I could not include it in my JIT, thus I had to use os/exec for invoking the compiler: GO 1 package main 2 3 import ( 4 "os" 5 "os/exec" 6 "fmt" 7 ) 8 9 func JIT () ( func () float64 , error ) { 10 generatedCode := `package main 11 func Main() float64 { 12 return 1E+00+1.2E+00-5E+00 13 } 14 ` 15 err := os . WriteFile ( "jit_output.go" , [] byte ( generatedCode ), 0777 ) 16 if err != nil { 17 return nil , err 18 } 19 cmd := exec . Command ( "go" , "build" , "-buildmode=plugin" , "-o" , "jit_output.so" , "jit_output.go" ) 20 err := cmd . Run () 21 if err != nil { 22 return nil , err 23 } 24 25 26 27 return nil , nil 28 } 29 30 func main () { 31 function , err := JIT () 32 if err != nil { 33 panic ( err ) 34 } 35 fmt . Println ( function ()) 36 } This first step calls the compiler and wants it to build a go plugin. Tip See go help build : TEXT 1 [...] 2 3 -buildmode mode 4 build mode to use. See 'go help buildmode' for more. 5 6 [...] And go help buildmode : TEXT 1 The 'go build' and 'go install' commands take a -buildmode argument which 2 indicates which kind of object file is to be built. Currently supported values 3 are: 4 5 [...] 6 7 -buildmode=plugin 8 Build the listed main packages, plus all packages that they 9 import, into a Go plugin. Packages not named main are ignored. 10 [...] Now a go plugin called jit_output.so sits in our project root. Go plugins Go features a package called plugin for loading and resolving go symbols of go plugins. There are some drawbacks such as missing portability due to no windows support and easily exploitable bugs in plugin loaders - since this is a POC JIT requiring the go compiler toolchain in the path I’m not even going to walk down the road of portability. Since we know the location of the plugin and the symbols in it our interactions with the plugin packages should be minimal - I need to open the plugin file, locate the Main function, cast it to func() float64 and return the result. GO 1 package main 2 3 import ( 4 "os" 5 "os/exec" 6 "fmt" 7 "plugin" 8 "errors" 9 ) 10 11 func JIT () ( func () float64 , error ) { 12 generatedCode := `package main 13 func Main() float64 { 14 return 1E+00+1.2E+00-5E+00 15 } 16 ` 17 err := os . WriteFile ( "jit_output.go" , [] byte ( generatedCode ), 0777 ) 18 if err != nil { 19 return nil , err 20 } 21 cmd := exec . Command ( "go" , "build" , "-buildmode=plugin" , "-o" , "jit_output.so" , "jit_output.go" ) 22 err := cmd . Run () 23 if err != nil { 24 return nil , err 25 } 26 27 plug , err := plugin . Open ( "jit_output.so" ) 28 if err != nil { 29 return nil , err 30 } 31 32 symbol , err := plug . Lookup ( "Main" ) 33 if err != nil { 34 fmt . Println ( sharedObjectPath ) 35 return nil , err 36 } 37 38 Main , ok := symbol .( func () float64 ) 39 if ! ok { 40 return nil , errors . New ( "Error while accessing jit compiled symbols" ) 41 } 42 43 return Main , nil 44 } 45 46 func main () { 47 function , err := JIT () 48 if err != nil { 49 panic ( err ) 50 } 51 fmt . Println ( function ()) 52 } What to do with this JIT My plan is to compare the three different approaches to evaluation techniques in a blog article in the next few weeks, expect some in depth benchmarks.