enzyme_go.go raw
1 // Bootstrap Go enzyme organ: tokenizes Go source code into typed elements.
2 //
3 // This is the hand-written template for the first enzyme organ. The organism
4 // will eventually learn to modify and improve it, but this provides the
5 // initial capability to digest Go source.
6 //
7 // Opcodes:
8 // 0x01 CanDigest(sample) -> 0x01 (yes) or 0x00 (no)
9 // 0x02 Digest(source) -> JSON array of {type, value} pairs
10 //
11 // Compile with:
12 // GOOS=wasip1 GOARCH=wasm go build -buildmode=c-shared -o enzyme_go.wasm .
13 package main
14
15 import (
16 "encoding/json"
17 "go/scanner"
18 "go/token"
19 "unsafe"
20 )
21
22 // Memory pool: bump allocator for guest-side allocations.
23 var pool [1 << 20]byte // 1MB
24 var poolOffset uint32
25
26 //go:wasmexport alloc
27 func alloc(size uint32) uint32 {
28 if poolOffset+size > uint32(len(pool)) {
29 poolOffset = 0
30 }
31 ptr := poolOffset
32 poolOffset += size
33 return uint32(uintptr(unsafe.Pointer(&pool[ptr])))
34 }
35
36 //go:wasmexport dealloc
37 func dealloc(ptr uint32, size uint32) {}
38
39 // Element is a typed lattice element produced by digestion.
40 type Element struct {
41 Type string `json:"type"`
42 Value string `json:"value"`
43 }
44
45 //go:wasmexport process
46 func process(inputPtr uint32, inputLen uint32) uint64 {
47 base := uint32(uintptr(unsafe.Pointer(&pool[0])))
48 offset := inputPtr - base
49
50 if offset+inputLen > uint32(len(pool)) {
51 return 0
52 }
53
54 input := pool[offset : offset+inputLen]
55 if len(input) == 0 {
56 return 0
57 }
58
59 opcode := input[0]
60 data := input[1:]
61
62 var result []byte
63
64 switch opcode {
65 case 0x01: // CanDigest: check if input looks like Go source.
66 result = []byte{0x00}
67 for i := 0; i+8 <= len(data); i++ {
68 if string(data[i:i+8]) == "package " {
69 result = []byte{0x01}
70 break
71 }
72 }
73
74 case 0x02: // Digest: tokenize Go source into elements.
75 elements := tokenize(data)
76 result, _ = json.Marshal(elements)
77 if result == nil {
78 result = []byte("[]")
79 }
80
81 default:
82 result = []byte("unknown opcode")
83 }
84
85 outPtr := alloc(uint32(len(result)))
86 outOffset := outPtr - base
87 copy(pool[outOffset:], result)
88
89 return uint64(outPtr)<<32 | uint64(len(result))
90 }
91
92 // tokenize uses go/scanner to break Go source into typed elements.
93 func tokenize(src []byte) []Element {
94 fset := token.NewFileSet()
95 file := fset.AddFile("input.go", fset.Base(), len(src))
96
97 var s scanner.Scanner
98 s.Init(file, src, nil, scanner.ScanComments)
99
100 var elements []Element
101
102 for {
103 _, tok, lit := s.Scan()
104 if tok == token.EOF {
105 break
106 }
107
108 elem := Element{
109 Type: classifyToken(tok),
110 Value: lit,
111 }
112 if lit == "" {
113 elem.Value = tok.String()
114 }
115
116 elements = append(elements, elem)
117 }
118
119 return elements
120 }
121
122 // classifyToken maps Go token types to lattice element types.
123 func classifyToken(tok token.Token) string {
124 switch {
125 case tok == token.IDENT:
126 return "ident"
127 case tok == token.INT || tok == token.FLOAT || tok == token.IMAG ||
128 tok == token.CHAR || tok == token.STRING:
129 return "literal"
130 case tok == token.COMMENT:
131 return "comment"
132 case tok == token.PACKAGE:
133 return "package"
134 case tok == token.IMPORT:
135 return "import"
136 case tok == token.FUNC:
137 return "func"
138 case tok == token.TYPE:
139 return "type"
140 case tok == token.STRUCT:
141 return "struct"
142 case tok == token.INTERFACE:
143 return "interface"
144 case tok == token.RETURN:
145 return "return"
146 case tok == token.IF:
147 return "if"
148 case tok == token.FOR:
149 return "for"
150 case tok == token.SWITCH:
151 return "switch"
152 case tok == token.SELECT:
153 return "select"
154 case tok == token.GO:
155 return "go"
156 case tok == token.CHAN:
157 return "chan"
158 case tok == token.VAR || tok == token.CONST:
159 return "var"
160 case tok.IsOperator():
161 return "op"
162 default:
163 return "punct"
164 }
165 }
166
167 func main() {}
168