1 files changed, 51 insertions, 25 deletions

diff --git a/nexgb/xgbgen/context.go b/nexgb/xgbgen/context.go
index a7a1d1d..67699cf 100644
--- a/nexgb/xgbgen/context.go
+++ b/nexgb/xgbgen/context.go
@@ -5,7 +5,6 @@ import (
 	"encoding/xml"
 	"fmt"
 	"log"
-	"strings"
 	"time"
 )
 
@@ -49,10 +48,10 @@ func (c *Context) Morph(xmlBytes []byte) {
 	parsedXml.Imports.Eval()
 
 	// Translate XML types to nice types
-	c.protocol = parsedXml.Translate()
+	c.protocol = parsedXml.Translate(nil)
 
 	// Start with Go header.
-	c.Putln("package xgb")
+	c.Putln("package %s", c.protocol.PkgName())
 	c.Putln("")
 	c.Putln("/*")
 	c.Putln("\tThis file was generated by %s.xml on %s.",
@@ -61,44 +60,53 @@ func (c *Context) Morph(xmlBytes []byte) {
 	c.Putln("*/")
 	c.Putln("")
 
-	// Write imports in comments
-	if len(c.protocol.Imports) > 0 {
-		c.Putln("// Imports are not necessary for XGB because everything is ")
-		c.Putln("// in one package. They are still listed here for reference.")
-		for _, imp := range c.protocol.Imports {
-			c.Putln("// import \"%s\"", imp.Name)
+	// Write imports. We always need to import at least xgb.
+	// We also need to import xproto if it's an extension.
+	c.Putln("import (")
+	c.Putln("\"github.com/BurntSushi/xgb\"")
+	c.Putln("")
+	if c.protocol.isExt() {
+		c.Putln("\"github.com/BurntSushi/xgb/xproto\"")
+	}
+	for _, imp := range c.protocol.Imports {
+		// We always import xproto, so skip it if it's explicitly imported
+		if imp.Name == "xproto" {
+			continue
 		}
-		c.Putln("")
+		c.Putln("\"github.com/BurntSushi/xgb/%s\"", imp.Name)
 	}
+	c.Putln(")")
+	c.Putln("")
 
 	// If this is an extension, create a function to initialize the extension
 	// before it can be used.
 	if c.protocol.isExt() {
-		name := strings.Title(c.protocol.Name) + "Init"
 		xname := c.protocol.ExtXName
 
-		c.Putln("// %s must be called before using the %s extension.",
-			name, xname)
-		c.Putln("func (c *Conn) %s() error {", name)
-		c.Putln("reply, err := c.QueryExtension(%d, \"%s\").Reply()",
+		c.Putln("// Init must be called before using the %s extension.",
+			xname)
+		c.Putln("func Init(c *xgb.Conn) error {")
+		c.Putln("reply, err := xproto.QueryExtension(c, %d, \"%s\").Reply()",
 			len(xname), xname)
 		c.Putln("switch {")
 		c.Putln("case err != nil:")
 		c.Putln("return err")
 		c.Putln("case !reply.Present:")
-		c.Putln("return errorf(\"No extension named %s could be found on "+
+		c.Putln("return xgb.Errorf(\"No extension named %s could be found on "+
 			"on the server.\")", xname)
 		c.Putln("}")
 		c.Putln("")
-		c.Putln("c.extLock.Lock()")
-		c.Putln("c.extensions[\"%s\"] = reply.MajorOpcode", xname)
-		c.Putln("for evNum, fun := range newExtEventFuncs[\"%s\"] {", xname)
-		c.Putln("newEventFuncs[int(reply.FirstEvent) + evNum] = fun")
+		c.Putln("xgb.ExtLock.Lock()")
+		c.Putln("c.Extensions[\"%s\"] = reply.MajorOpcode", xname)
+		c.Putln("for evNum, fun := range xgb.NewExtEventFuncs[\"%s\"] {",
+			xname)
+		c.Putln("xgb.NewEventFuncs[int(reply.FirstEvent) + evNum] = fun")
 		c.Putln("}")
-		c.Putln("for errNum, fun := range newExtErrorFuncs[\"%s\"] {", xname)
-		c.Putln("newErrorFuncs[int(reply.FirstError) + errNum] = fun")
+		c.Putln("for errNum, fun := range xgb.NewExtErrorFuncs[\"%s\"] {",
+			xname)
+		c.Putln("xgb.NewErrorFuncs[int(reply.FirstError) + errNum] = fun")
 		c.Putln("}")
-		c.Putln("c.extLock.Unlock()")
+		c.Putln("xgb.ExtLock.Unlock()")
 		c.Putln("")
 		c.Putln("return nil")
 		c.Putln("}")
@@ -107,8 +115,26 @@ func (c *Context) Morph(xmlBytes []byte) {
 		// Make sure newExtEventFuncs["EXT_NAME"] map is initialized.
 		// Same deal for newExtErrorFuncs["EXT_NAME"]
 		c.Putln("func init() {")
-		c.Putln("newExtEventFuncs[\"%s\"] = make(map[int]newEventFun)", xname)
-		c.Putln("newExtErrorFuncs[\"%s\"] = make(map[int]newErrorFun)", xname)
+		c.Putln("xgb.NewExtEventFuncs[\"%s\"] = make(map[int]xgb.NewEventFun)",
+			xname)
+		c.Putln("xgb.NewExtErrorFuncs[\"%s\"] = make(map[int]xgb.NewErrorFun)",
+			xname)
+		c.Putln("}")
+		c.Putln("")
+	} else {
+		// In the xproto package, we must provide a Setup function that uses
+		// SetupBytes in xgb.Conn to return a SetupInfo structure.
+		c.Putln("// Setup parses the setup bytes retrieved when")
+		c.Putln("// connecting into a SetupInfo struct.")
+		c.Putln("func Setup(c *xgb.Conn) *SetupInfo {")
+		c.Putln("setup := new(SetupInfo)")
+		c.Putln("SetupInfoRead(c.SetupBytes, setup)")
+		c.Putln("return setup")
+		c.Putln("}")
+		c.Putln("")
+		c.Putln("// DefaultScreen gets the default screen info from SetupInfo.")
+		c.Putln("func (s *SetupInfo) DefaultScreen(c *xgb.Conn) *ScreenInfo {")
+		c.Putln("return &s.Roots[c.DefaultScreen]")
 		c.Putln("}")
 		c.Putln("")
 	}