// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

// Stub package for testing ssa compiler backend.  Will eventually
// be deleted when ssa is called directly from the main compiler.

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"

	"cmd/internal/ssa/types"

	"cmd/internal/ssa"
)

// testing harness which runs the compiler using an IR read from a file
func main() {
	flag.Parse()
	file := flag.Arg(0)
	r, err := os.Open(file)
	if err != nil {
		panic(err)
	}
	f := buildFunc(readFunc(r))
	ssa.Compile(f)
	// TODO: output f
}

// readFunc reads the intermediate representation generated by the
// compiler frontend and returns it as a list of sexpressions.
func readFunc(r io.Reader) []sexpr {
	var lines []sexpr
	s := bufio.NewScanner(r)
	for s.Scan() {
		line := s.Text()
		e := parseSexpr(strings.Trim(line, " "))

		if !e.compound {
			panic("bad stmt: " + line)
		}
		if e.parts[0].compound {
			panic("bad op: " + line)
		}
		lines = append(lines, e)
	}
	return lines
}

// buildFunc converts from the 6g IR dump format to the internal
// form.  Builds SSA and all that.
func buildFunc(lines []sexpr) *ssa.Func {
	f := new(ssa.Func)

	// We construct SSA using an algorithm similar to
	// Brau, Buchwald, Hack, Leißa, Mallon, and Zwinkau
	// http://pp.info.uni-karlsruhe.de/uploads/publikationen/braun13cc.pdf

	// allocate starting block
	f.Entry = f.NewBlock(ssa.BlockPlain)
	// TODO: all args.  Make a struct containing args/returnvals, declare
	// an FP which contains a pointer to that struct.

	var exit *ssa.Block // all returns (if any) branch to here  TODO: defers & panics?

	// add a block for each label
	// Also a few other preprocessing steps, all in one pass.
	labels := map[string]*ssa.Block{}
	types := map[string]ssa.Type{}
	callFallthrough := map[int]*ssa.Block{}
	for i, e := range lines {
		switch e.parts[0].name {
		case "LABEL":
			labels[e.parts[1].name] = f.NewBlock(ssa.BlockPlain)
		case "NAME":
			f.Name = e.parts[1].name
		case "RETURN":
			if exit == nil {
				exit = f.NewBlock(ssa.BlockExit)
			}
		case "TYPE":
			types[e.parts[1].name] = parseSexprType(e.parts[2])
		case "CALL":
			// allocate a new block for fallthrough
			callFallthrough[i] = f.NewBlock(ssa.BlockPlain)
			if exit == nil {
				exit = f.NewBlock(ssa.BlockExit)
			}
		}
	}

	// map from block id to sexprs in that block
	blocklines := make([][]sexpr, f.NumBlocks())

	// Add sexprs to the correct block.  Add edges between blocks.
	b := f.Entry
	var i int
	for j, e := range lines {
		if b == nil && e.parts[0].name != "LABEL" {
			// dead code (e.g. return in "if" branch makes the "goto end" statement dead)
			continue
		}
		switch e.parts[0].name {
		case "IF":
			if b.Kind != ssa.BlockPlain {
				panic("bad b state")
			}
			b.Kind = ssa.BlockIf
			edge(b, labels[e.parts[2].name])
			edge(b, labels[e.parts[3].name])
			blocklines[b.ID] = lines[i : j+1]
			b = nil
		case "GOTO":
			edge(b, labels[e.parts[1].name])
			blocklines[b.ID] = lines[i:j]
			b = nil
		case "LABEL":
			b = labels[e.parts[1].name]
			i = j + 1
		case "RETURN":
			if b.Kind != ssa.BlockPlain {
				panic("bad b state")
			}
			edge(b, exit)
			blocklines[b.ID] = lines[i:j]
			b = nil
		case "CALL":
			if b.Kind != ssa.BlockPlain {
				panic("bad b state")
			}
			b.Kind = ssa.BlockCall
			c := callFallthrough[j]
			edge(b, c)
			edge(b, exit)
			blocklines[b.ID] = lines[i : j+1]
			b = c
			i = j + 1
		}
		// note that we don't keep goto/label/return sexprs
	}
	if b != nil {
		panic("control flow falls off end of function")
	}

	// Read types for each variable
	// Number the variables densely
	varids := map[string]int{} // map from variable name to id
	var varnames []string      // map from id to variable name
	var vartypes []ssa.Type    // map from variable id to type
	for _, e := range lines {
		if e.parts[0].name != "DCL" {
			continue
		}
		name := e.parts[1].name
		if _, ok := varids[name]; ok {
			continue
		}
		id := len(varids)
		if id == 1<<31-1 {
			panic("too many variables")
		}
		fmt.Printf("var %d = %s\n", id, name)
		varids[name] = id
		varnames = append(varnames, name)
		vartypes = append(vartypes, types[e.parts[2].name])
	}
	memID := len(varids)
	fmt.Printf("var %d = .mem\n", memID)
	varids[".mem"] = memID // TODO: need .mem here?
	varnames = append(varnames, ".mem")
	vartypes = append(vartypes, ssa.TypeMem)

	// map from variable ID to current Value of that variable
	curBlock := NewSparseMap(len(varids))

	var state ssaFuncState
	state.types = types
	state.varids = varids
	state.varnames = varnames
	state.vartypes = vartypes
	state.curBlock = curBlock
	state.done = make([]bool, f.NumBlocks())
	state.defs = map[blockvar]*ssa.Value{}
	state.memID = memID

	// Convert each block to ssa
	// TODO: order blocks for maximum happiness - we want to process
	// all the predecessors of a block before processing the block itself,
	// if at all possible.
	for _, b := range f.Blocks {
		fmt.Printf("processing block %d\n", b.ID)
		curBlock.Clear()
		for _, e := range blocklines[b.ID] {
			switch e.parts[0].name {
			case "AS":
				if e.parts[1].compound {
					// store expression
					lhs := genExpr(&state, b, e.parts[1])
					rhs := genExpr(&state, b, e.parts[2])
					mem := genVar(&state, b, memID)
					v := b.NewValue(ssa.OpStore, ssa.TypeMem, nil)
					v.AddArg(lhs)
					v.AddArg(rhs)
					v.AddArg(mem)
					curBlock.Put(memID, v)
				} else {
					// variable assignment
					v := genExpr(&state, b, e.parts[2])
					curBlock.Put(varids[e.parts[1].name], v)
				}
			case "DCL":
				// nothing to do
			case "IF":
				b.Control = genExpr(&state, b, e.parts[1])
			case "CALL":
				// only direct call for now - indirect call takes addr value as well
				v := b.NewValue(ssa.OpStaticCall, ssa.TypeMem, e.parts[1].name)
				v.AddArg(genVar(&state, b, memID))
				curBlock.Put(memID, v)
				b.Control = v
			}
		}
		// link up thunks to their actual values
		for _, v := range b.Values {
			if v.Op != ssa.OpThunk {
				continue
			}
			varid := v.Aux.(int)
			w := genVar(&state, b, varid)
			v.Op = ssa.OpCopy
			v.Aux = nil
			v.AddArg(w)
		}

		// record final values at the end of the block
		for _, e := range curBlock.Contents() {
			state.defs[blockvar{b.ID, e.Key}] = e.Val
			// TODO: somehow avoid storing dead values to this map.
		}
		curBlock.Clear()
		state.done[b.ID] = true
	}

	// the final store value is returned
	if exit != nil {
		exit.Control = genVar(&state, exit, memID)
	}

	return f
}

func edge(a, b *ssa.Block) {
	a.Succs = append(a.Succs, b)
	b.Preds = append(b.Preds, a)
}

func genVar(state *ssaFuncState, b *ssa.Block, id int) *ssa.Value {
	// look up variable
	v := state.curBlock.Get(id)
	if v != nil {
		// variable was defined previously in this block
		// (or we memoized the result)
		return v
	}

	// Variable comes in from outside of basic block.
	v = lookupVarIncoming(state, b, id)

	// memoize result so future callers will not look it up again
	state.curBlock.Put(id, v)
	return v
}

func genExpr(state *ssaFuncState, b *ssa.Block, e sexpr) *ssa.Value {
	if !e.compound {
		return genVar(state, b, state.varids[e.name])
	}
	switch e.parts[0].name {
	case "ADD":
		x := genExpr(state, b, e.parts[1])
		y := genExpr(state, b, e.parts[2])
		v := b.NewValue(ssa.OpAdd, x.Type, nil)
		v.AddArg(x)
		v.AddArg(y)
		return v
	case "SUB":
		x := genExpr(state, b, e.parts[1])
		y := genExpr(state, b, e.parts[2])
		v := b.NewValue(ssa.OpSub, x.Type, nil)
		v.AddArg(x)
		v.AddArg(y)
		return v
	case "CINT":
		c, err := strconv.ParseInt(e.parts[1].name, 10, 64)
		if err != nil {
			panic("bad cint value")
		}
		return b.Func.ConstInt(c)
	case "LT":
		x := genExpr(state, b, e.parts[1])
		y := genExpr(state, b, e.parts[2])
		v := b.NewValue(ssa.OpLess, ssa.TypeBool, nil)
		v.AddArg(x)
		v.AddArg(y)
		return v
	case "FP":
		typ := state.types[e.parts[1].name]
		offset, err := strconv.ParseInt(e.parts[2].name, 10, 64)
		if err != nil {
			panic(err)
		}
		v := b.NewValue(ssa.OpFPAddr, types.NewPointer(typ), offset)
		return v
	case "SP":
		typ := state.types[e.parts[1].name]
		offset, err := strconv.ParseInt(e.parts[2].name, 10, 64)
		if err != nil {
			panic(err)
		}
		v := b.NewValue(ssa.OpSPAddr, types.NewPointer(typ), offset)
		return v
	case "LOAD":
		p := genExpr(state, b, e.parts[1])
		v := b.NewValue(ssa.OpLoad, p.Type.(*types.Pointer).Elem(), nil)
		v.AddArg(p)
		v.AddArg(genVar(state, b, state.memID))
		return v
	default:
		fmt.Println(e.parts[0].name)
		panic("unknown op")
	}
}

// map key combining block id and variable id
type blockvar struct {
	bid   ssa.ID
	varid int
}

type ssaFuncState struct {
	types    map[string]ssa.Type
	varnames []string
	varids   map[string]int
	vartypes []ssa.Type
	curBlock *SparseMap              // value of each variable in block we're working on
	defs     map[blockvar]*ssa.Value // values for variables at the end of blocks
	done     []bool
	memID    int
}

// Find the value of the variable with the given id leaving block b.
func lookupVarOutgoing(state *ssaFuncState, b *ssa.Block, id int) *ssa.Value {
	fmt.Printf("lookupOutgoing var=%d block=%d\n", id, b.ID)
	v := state.defs[blockvar{b.ID, id}]
	if v != nil {
		return v
	}
	if state.done[b.ID] {
		// The variable was not defined in this block, and we haven't
		// memoized the answer yet.  Look it up recursively.  This might
		// cause infinite recursion, so add a copy first.
		v = b.NewValue(ssa.OpCopy, state.vartypes[id], nil)
		state.defs[blockvar{b.ID, id}] = v
		v.AddArg(lookupVarIncoming(state, b, id))
		return v
	}
	// We don't know about defined variables in this block (yet).
	// Make a thunk for this variable.
	fmt.Printf("making thunk for var=%d in block=%d\n", id, b.ID)
	v = b.NewValue(ssa.OpThunk, state.vartypes[id], id)

	// memoize result
	state.defs[blockvar{b.ID, id}] = v
	return v
}

// Find the Value of the variable coming into block b.
func lookupVarIncoming(state *ssaFuncState, b *ssa.Block, id int) *ssa.Value {
	fmt.Printf("lookupIncoming var=%d block=%d\n", id, b.ID)
	var v *ssa.Value
	switch len(b.Preds) {
	case 0:
		// TODO: handle function args some other way (assignments in starting block?)
		// TODO: error if variable isn't a function arg (including mem input)
		v = b.NewValue(ssa.OpArg, state.vartypes[id], state.varnames[id])
	case 1:
		v = lookupVarOutgoing(state, b.Preds[0], id)
	default:
		v = b.NewValue(ssa.OpCopy, state.vartypes[id], nil)

		args := make([]*ssa.Value, len(b.Preds))
		for i, p := range b.Preds {
			args[i] = lookupVarOutgoing(state, p, id)
		}

		// if <=1 value that isn't this variable's thunk, don't make phi
		v.Op = ssa.OpPhi
		v.AddArgs(args...) // note: order corresponding to b.Pred
	}
	return v
}

func parseSexprType(e sexpr) ssa.Type {
	if !e.compound {
		switch e.name {
		case "int":
			return ssa.TypeInt
		default:
			fmt.Println(e.name)
			panic("unknown type")
		}
	}
	if e.parts[0].name == "FUNC" {
		// TODO: receiver?  Already folded into args?  Variadic?
		var args, rets []*types.Var
		for _, s := range e.parts[1].parts {
			t := parseSexprType(s)
			args = append(args, types.NewParam(0, nil, "noname", t))
		}
		for _, s := range e.parts[2].parts {
			t := parseSexprType(s)
			rets = append(rets, types.NewParam(0, nil, "noname", t))
		}
		sig := types.NewSignature(nil, nil, types.NewTuple(args...), types.NewTuple(rets...), false)
		return ssa.Type(sig)
	}
	// TODO: array/struct/...
	panic("compound type")
}