compileBlock: support pyModule init

cl/compile.go

@@ -216,7 +216,7 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) llssa.Func
 	if ftype != goFunc {
 		if ftype == pyFunc {
 			// TODO(xsw): pyMod == ""
-			fn := "__llgo_py." + p.pyMod + "." + name
+			fn := pysymPrefix + p.pyMod + "." + name
 			pkg.NewVar(fn, types.Typ[types.Int], llssa.InC)
 		}
 		return nil
@@ -250,6 +250,7 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) llssa.Func
 	}
 	if nblk := len(f.Blocks); nblk > 0 {
 		fn.MakeBlocks(nblk) // to set fn.HasBody() = true
+		isPyMod := p.pyMod != ""
 		p.inits = append(p.inits, func() {
 			p.fn = fn
 			defer func() {
@@ -269,7 +270,9 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function) llssa.Func
 				off[i] = p.compilePhis(b, block)
 			}
 			for i, block := range f.Blocks {
-				p.compileBlock(b, block, off[i], i == 0 && name == "main")
+				doInit := (i == 0 && name == "main")
+				doPyModInit := (isPyMod && i == 1 && f.Name() == "init" && sig.Recv() == nil)
+				p.compileBlock(b, block, off[i], doInit, doPyModInit)
 			}
 			for _, phi := range p.phis {
 				phi()
@@ -314,25 +317,46 @@ func (p *context) funcOf(fn *ssa.Function) (ret llssa.Function, ftype int) {
 	return
 }
 
-func (p *context) compileBlock(b llssa.Builder, block *ssa.BasicBlock, n int, doInit bool) llssa.BasicBlock {
-	ret := p.fn.Block(block.Index)
+func (p *context) compileBlock(b llssa.Builder, block *ssa.BasicBlock, n int, doInit, pyModInit bool) llssa.BasicBlock {
+	var last int
+	var instrs []ssa.Instruction
+	var ret = p.fn.Block(block.Index)
 	b.SetBlock(ret)
-	if doInit {
-		prog := p.prog
-		pkg := p.pkg
-		fn := p.fn
-		argc := pkg.NewVar("__llgo_argc", types.NewPointer(types.Typ[types.Int32]), llssa.InC)
-		argv := pkg.NewVar("__llgo_argv", types.NewPointer(argvTy), llssa.InC)
-		argc.Init(prog.Null(argc.Type))
-		argv.Init(prog.Null(argv.Type))
-		b.Store(argc.Expr, fn.Param(0))
-		b.Store(argv.Expr, fn.Param(1))
-		callRuntimeInit(b, pkg)
-		b.Call(pkg.FuncOf("main.init").Expr)
+	if pyModInit {
+		last = len(block.Instrs) - 1
+		instrs = block.Instrs[n:last]
+	} else {
+		instrs = block.Instrs[n:]
+		if doInit {
+			prog := p.prog
+			pkg := p.pkg
+			fn := p.fn
+			argc := pkg.NewVar("__llgo_argc", types.NewPointer(types.Typ[types.Int32]), llssa.InC)
+			argv := pkg.NewVar("__llgo_argv", types.NewPointer(argvTy), llssa.InC)
+			argc.Init(prog.Null(argc.Type))
+			argv.Init(prog.Null(argv.Type))
+			b.Store(argc.Expr, fn.Param(0))
+			b.Store(argv.Expr, fn.Param(1))
+			callRuntimeInit(b, pkg)
+			b.Call(pkg.FuncOf("main.init").Expr)
+		}
 	}
-	for _, instr := range block.Instrs[n:] {
+	for _, instr := range instrs {
 		p.compileInstr(b, instr)
 	}
+	if pyModInit {
+		jump := block.Instrs[last].(*ssa.Jump)
+		jumpTo := p.jumpTo(jump)
+		modName := pysymPrefix + p.pyMod
+		modPtr := p.pkg.NewPyModVar(modName).Expr
+		mod := b.Load(modPtr)
+		cond := b.BinOp(token.NEQ, mod, b.Prog.Null(mod.Type))
+		newBlk := p.fn.MakeBlock()
+		b.If(cond, jumpTo, newBlk)
+		b.SetBlock(newBlk)
+		// TODO(xsw): pyModInit
+		b.Jump(jumpTo)
+	}
 	return ret
 }
 
@@ -657,6 +681,12 @@ func (p *context) compileInstrOrValue(b llssa.Builder, iv instrOrValue, asValue
 	return ret
 }
 
+func (p *context) jumpTo(v *ssa.Jump) llssa.BasicBlock {
+	fn := p.fn
+	succs := v.Block().Succs
+	return fn.Block(succs[0].Index)
+}
+
 func (p *context) compileInstr(b llssa.Builder, instr ssa.Instruction) {
 	if iv, ok := instr.(instrOrValue); ok {
 		p.compileInstrOrValue(b, iv, false)
@@ -680,9 +710,7 @@ func (p *context) compileInstr(b llssa.Builder, instr ssa.Instruction) {
 		val := p.compileValue(b, v.Val)
 		b.Store(ptr, val)
 	case *ssa.Jump:
-		fn := p.fn
-		succs := v.Block().Succs
-		jmpb := fn.Block(succs[0].Index)
+		jmpb := p.jumpTo(v)
 		b.Jump(jmpb)
 	case *ssa.Return:
 		var results []llssa.Expr

cl/import.go

@@ -402,6 +402,10 @@ func pkgKindByPath(pkgPath string) int {
 
 // -----------------------------------------------------------------------------
 
+const (
+	pysymPrefix = "__llgo_py."
+)
+
 func (p *context) initPyModule() {
 	if kind, mod := pkgKindByScope(p.goTyps.Scope()); kind == PkgPyModule {
 		p.pyMod = mod

ssa/decl.go

@@ -212,19 +212,29 @@ func (p Function) MakeBody(nblk int) Builder {
 
 // MakeBlocks creates nblk basic blocks for the function.
 func (p Function) MakeBlocks(nblk int) []BasicBlock {
-	if p.blks == nil {
+	n := len(p.blks)
+	if n == 0 {
 		p.blks = make([]BasicBlock, 0, nblk)
 	}
-	n := len(p.blks)
-	f := p.impl
 	for i := 0; i < nblk; i++ {
-		label := "_llgo_" + strconv.Itoa(i)
-		blk := llvm.AddBasicBlock(f, label)
-		p.blks = append(p.blks, &aBasicBlock{blk, p, n + i})
+		p.addBlock(n + i)
 	}
 	return p.blks[n:]
 }
 
+func (p Function) addBlock(idx int) BasicBlock {
+	label := "_llgo_" + strconv.Itoa(idx)
+	blk := llvm.AddBasicBlock(p.impl, label)
+	ret := &aBasicBlock{blk, p, idx}
+	p.blks = append(p.blks, ret)
+	return ret
+}
+
+// MakeBlock creates a new basic block for the function.
+func (p Function) MakeBlock() BasicBlock {
+	return p.addBlock(len(p.blks))
+}
+
 // Block returns the ith basic block of the function.
 func (p Function) Block(idx int) BasicBlock {
 	return p.blks[idx]

ssa/package.go

@@ -327,6 +327,13 @@ func (p Package) NewConst(name string, val constant.Value) NamedConst {
 }
 */
 
+// NewPyModVar creates a new global variable for a Python module.
+func (p Package) NewPyModVar(name string) Global {
+	ret := p.NewVar(name, types.NewPointer(types.Typ[types.Int]), InC)
+	ret.impl.SetLinkage(llvm.LinkOnceAnyLinkage)
+	return ret
+}
+
 // NewVar creates a new global variable.
 func (p Package) NewVar(name string, typ types.Type, bg Background) Global {
 	t := p.Prog.Type(typ, bg)