/*

 * This file compiles an abstract syntax tree (AST) into Python bytecode.

 *

 * The primary entry point is PyAST_Compile(), which returns a

 * PyCodeObject.  The compiler makes several passes to build the code

 * object:

 *   1. Checks for future statements.  See future.c

 *   2. Builds a symbol table.	See symtable.c.

 *   3. Generate code for basic blocks.  See compiler_mod() in this file.

 *   4. Assemble the basic blocks into final code.  See assemble() in

 *	this file.	 

 *   5. Optimize the byte code (peephole optimizations).  See peephole.c

 *

 * Note that compiler_mod() suggests module, but the module ast type

 * (mod_ty) has cases for expressions and interactive statements.

 *

 * CAUTION: The VISIT_* macros abort the current function when they

 * encounter a problem. So don't invoke them when there is memory

 * which needs to be released. Code blocks are OK, as the compiler

 * structure takes care of releasing those.  Use the arena to manage

 * objects.

 */

#include "Python.h"

#include "Python-ast.h"

#include "node.h"

#include "pyarena.h"

#include "ast.h"

#include "code.h"

#include "compile.h"

#include "symtable.h"

#include "opcode.h"

int Py_OptimizeFlag = 0;

#define DEFAULT_BLOCK_SIZE 16

#define DEFAULT_BLOCKS 8

#define DEFAULT_CODE_SIZE 128

#define DEFAULT_LNOTAB_SIZE 16

struct instr {

	unsigned i_jabs : 1;

	unsigned i_jrel : 1;

	unsigned i_hasarg : 1;

	unsigned char i_opcode;

	int i_oparg;

	struct basicblock_ *i_target; /* target block (if jump instruction) */

	int i_lineno;

};

typedef struct basicblock_ {

    /* Each basicblock in a compilation unit is linked via b_list in the

       reverse order that the block are allocated.  b_list points to the next

       block, not to be confused with b_next, which is next by control flow. */

	struct basicblock_ *b_list;

	/* number of instructions used */

	int b_iused;

	/* length of instruction array (b_instr) */

	int b_ialloc;

	/* pointer to an array of instructions, initially NULL */

	struct instr *b_instr;

	/* If b_next is non-NULL, it is a pointer to the next

	   block reached by normal control flow. */

	struct basicblock_ *b_next;

	/* b_seen is used to perform a DFS of basicblocks. */

	unsigned b_seen : 1;

	/* b_return is true if a RETURN_VALUE opcode is inserted. */

	unsigned b_return : 1;

	/* depth of stack upon entry of block, computed by stackdepth() */

	int b_startdepth;

	/* instruction offset for block, computed by assemble_jump_offsets() */

	int b_offset;

} basicblock;

/* fblockinfo tracks the current frame block.

A frame block is used to handle loops, try/except, and try/finally.

It's called a frame block to distinguish it from a basic block in the

compiler IR.

*/

enum fblocktype { LOOP, EXCEPT, FINALLY_TRY, FINALLY_END };

struct fblockinfo {

	enum fblocktype fb_type;

	basicblock *fb_block;

};

/* The following items change on entry and exit of code blocks.

   They must be saved and restored when returning to a block.

*/

struct compiler_unit {

	PySTEntryObject *u_ste;

	PyObject *u_name;

	/* The following fields are dicts that map objects to

	   the index of them in co_XXX.	 The index is used as

	   the argument for opcodes that refer to those collections.

	*/

	PyObject *u_consts;    /* all constants */

	PyObject *u_names;     /* all names */

	PyObject *u_varnames;  /* local variables */

	PyObject *u_cellvars;  /* cell variables */

	PyObject *u_freevars;  /* free variables */

	PyObject *u_private;	/* for private name mangling */

	int u_argcount;	   /* number of arguments for block */ 

	/* Pointer to the most recently allocated block.  By following b_list

	   members, you can reach all early allocated blocks. */

	basicblock *u_blocks;

	basicblock *u_curblock; /* pointer to current block */

	int u_tmpname;		/* temporary variables for list comps */

	int u_nfblocks;

	struct fblockinfo u_fblock[CO_MAXBLOCKS];

	int u_firstlineno; /* the first lineno of the block */

	int u_lineno;	   /* the lineno for the current stmt */

	bool u_lineno_set; /* boolean to indicate whether instr

			      has been generated with current lineno */

};

/* This struct captures the global state of a compilation.  

The u pointer points to the current compilation unit, while units

for enclosing blocks are stored in c_stack.	The u and c_stack are

managed by compiler_enter_scope() and compiler_exit_scope().

*/

struct compiler {

	const char *c_filename;

	struct symtable *c_st;

	PyFutureFeatures *c_future; /* pointer to module's __future__ */

	PyCompilerFlags *c_flags;

	int c_interactive;	 /* true if in interactive mode */

	int c_nestlevel;

	struct compiler_unit *u; /* compiler state for current block */

	PyObject *c_stack;	 /* Python list holding compiler_unit ptrs */

	char *c_encoding;	 /* source encoding (a borrowed reference) */

	PyArena *c_arena;	 /* pointer to memory allocation arena */

};

static int compiler_enter_scope(struct compiler *, identifier, void *, int);

static void compiler_free(struct compiler *);

static basicblock *compiler_new_block(struct compiler *);

static int compiler_next_instr(struct compiler *, basicblock *);

static int compiler_addop(struct compiler *, int);

static int compiler_addop_o(struct compiler *, int, PyObject *, PyObject *);

static int compiler_addop_i(struct compiler *, int, int);

static int compiler_addop_j(struct compiler *, int, basicblock *, int);

static basicblock *compiler_use_new_block(struct compiler *);

static int compiler_error(struct compiler *, const char *);

static int compiler_nameop(struct compiler *, identifier, expr_context_ty);

static PyCodeObject *compiler_mod(struct compiler *, mod_ty);

static int compiler_visit_stmt(struct compiler *, stmt_ty);

static int compiler_visit_keyword(struct compiler *, keyword_ty);

static int compiler_visit_expr(struct compiler *, expr_ty);

static int compiler_augassign(struct compiler *, stmt_ty);

static int compiler_visit_slice(struct compiler *, slice_ty,

				expr_context_ty);

static int compiler_push_fblock(struct compiler *, enum fblocktype,

				basicblock *);

static void compiler_pop_fblock(struct compiler *, enum fblocktype,

				basicblock *);

/* Returns true if there is a loop on the fblock stack. */

static int compiler_in_loop(struct compiler *);

static int inplace_binop(struct compiler *, operator_ty);

static int expr_constant(expr_ty e);

static int compiler_with(struct compiler *, stmt_ty);

static PyCodeObject *assemble(struct compiler *, int addNone);

static PyObject *__doc__;

PyObject *

_Py_Mangle(PyObject *privateobj, PyObject *ident)

{

	/* Name mangling: __private becomes _classname__private.

	   This is independent from how the name is used. */

	const char *p, *name = PyString_AsString(ident);

	char *buffer;

	size_t nlen, plen;

	if (privateobj == NULL || !PyString_Check(privateobj) ||

	    name == NULL || name[0] != '_' || name[1] != '_') {

		Py_INCREF(ident);

		return ident;

	}

	p = PyString_AsString(privateobj);

	nlen = strlen(name);

	/* Don't mangle __id__ or names with dots.

	   The only time a name with a dot can occur is when

	   we are compiling an import statement that has a 

	   package name.

	   TODO(jhylton): Decide whether we want to support

	   mangling of the module name, e.g. __M.X.

	*/

	if ((name[nlen-1] == '_' && name[nlen-2] == '_') 

	    || strchr(name, '.')) {

		Py_INCREF(ident);

		return ident; /* Don't mangle __whatever__ */

	}

	/* Strip leading underscores from class name */

	while (*p == '_')

		p++;

	if (*p == '\0') {

		Py_INCREF(ident);

		return ident; /* Don't mangle if class is just underscores */

	}

	plen = strlen(p);

	assert(1 <= PY_SSIZE_T_MAX - nlen);

	assert(1 + nlen <= PY_SSIZE_T_MAX - plen);

	ident = PyString_FromStringAndSize(NULL, 1 + nlen + plen);

	if (!ident)

		return 0;

	/* ident = "_" + p[:plen] + name # i.e. 1+plen+nlen bytes */

	buffer = PyString_AS_STRING(ident);

	buffer[0] = '_';

	strncpy(buffer+1, p, plen);

	strcpy(buffer+1+plen, name);

	return ident;

}

static int

compiler_init(struct compiler *c)

{

	memset(c, 0, sizeof(struct compiler));

	c->c_stack = PyList_New(0);

	if (!c->c_stack)

		return 0;

	return 1;

}

PyCodeObject *

PyAST_Compile(mod_ty mod, const char *filename, PyCompilerFlags *flags,

	      PyArena *arena)

{

	struct compiler c;

	PyCodeObject *co = NULL;

	PyCompilerFlags local_flags;

	int merged;

	if (!__doc__) {

		__doc__ = PyString_InternFromString("__doc__");

		if (!__doc__)

			return NULL;

	}

	if (!compiler_init(&c))

		return NULL;

	c.c_filename = filename;

	c.c_arena = arena;

	c.c_future = PyFuture_FromAST(mod, filename);

	if (c.c_future == NULL)

		goto finally;

	if (!flags) {

		local_flags.cf_flags = 0;

		flags = &local_flags;

	}

	merged = c.c_future->ff_features | flags->cf_flags;

	c.c_future->ff_features = merged;

	flags->cf_flags = merged;

	c.c_flags = flags;

	c.c_nestlevel = 0;

	c.c_st = PySymtable_Build(mod, filename, c.c_future);

	if (c.c_st == NULL) {

		if (!PyErr_Occurred())

			PyErr_SetString(PyExc_SystemError, "no symtable");

		goto finally;

	}

	/* XXX initialize to NULL for now, need to handle */

	c.c_encoding = NULL;

	co = compiler_mod(&c, mod);

 finally:

	compiler_free(&c);

	assert(co || PyErr_Occurred());

	return co;

}

PyCodeObject *

PyNode_Compile(struct _node *n, const char *filename)

{

	PyCodeObject *co = NULL;

	mod_ty mod;

	PyArena *arena = PyArena_New();

	if (!arena)

		return NULL;

	mod = PyAST_FromNode(n, NULL, filename, arena);

	if (mod)

		co = PyAST_Compile(mod, filename, NULL, arena);

	PyArena_Free(arena);

	return co;

}

static void

compiler_free(struct compiler *c)

{

	if (c->c_st)

		PySymtable_Free(c->c_st);

	if (c->c_future)

		PyObject_Free(c->c_future);

	Py_DECREF(c->c_stack);

}

static PyObject *

list2dict(PyObject *list)

{

	Py_ssize_t i, n;

	PyObject *v, *k;

	PyObject *dict = PyDict_New();

	if (!dict) return NULL;

	n = PyList_Size(list);

	for (i = 0; i < n; i++) {

		v = PyInt_FromLong(i);

		if (!v) {

			Py_DECREF(dict);

			return NULL;

		}

		k = PyList_GET_ITEM(list, i);

		k = PyTuple_Pack(2, k, k->ob_type);

		if (k == NULL || PyDict_SetItem(dict, k, v) < 0) {

			Py_XDECREF(k);

			Py_DECREF(v);

			Py_DECREF(dict);

			return NULL;

		}

		Py_DECREF(k);

		Py_DECREF(v);

	}

	return dict;

}

/* Return new dict containing names from src that match scope(s).

src is a symbol table dictionary.  If the scope of a name matches

either scope_type or flag is set, insert it into the new dict.	The

values are integers, starting at offset and increasing by one for

each key.

*/

static PyObject *

dictbytype(PyObject *src, int scope_type, int flag, int offset)

{

	Py_ssize_t pos = 0, i = offset, scope;

	PyObject *k, *v, *dest = PyDict_New();

	assert(offset >= 0);

	if (dest == NULL)

		return NULL;

	while (PyDict_Next(src, &pos, &k, &v)) {

		/* XXX this should probably be a macro in symtable.h */

		assert(PyInt_Check(v));

		scope = (PyInt_AS_LONG(v) >> SCOPE_OFF) & SCOPE_MASK;

		if (scope == scope_type || PyInt_AS_LONG(v) & flag) {

			PyObject *tuple, *item = PyInt_FromLong(i);

			if (item == NULL) {

				Py_DECREF(dest);

				return NULL;

			}

			i++;

			tuple = PyTuple_Pack(2, k, k->ob_type);

			if (!tuple || PyDict_SetItem(dest, tuple, item) < 0) {

				Py_DECREF(item);

				Py_DECREF(dest);

				Py_XDECREF(tuple);

				return NULL;

			}

			Py_DECREF(item);

			Py_DECREF(tuple);

		}

	}

	return dest;

}

static void

compiler_unit_check(struct compiler_unit *u)

{

	basicblock *block;

	for (block = u->u_blocks; block != NULL; block = block->b_list) {

		assert((void *)block != (void *)0xcbcbcbcb);

		assert((void *)block != (void *)0xfbfbfbfb);

		assert((void *)block != (void *)0xdbdbdbdb);

		if (block->b_instr != NULL) {

			assert(block->b_ialloc > 0);

			assert(block->b_iused > 0);

			assert(block->b_ialloc >= block->b_iused);

		}

		else {

			assert (block->b_iused == 0);

			assert (block->b_ialloc == 0);

		}

	}

}

static void

compiler_unit_free(struct compiler_unit *u)

{

	basicblock *b, *next;

	compiler_unit_check(u);

	b = u->u_blocks;

	while (b != NULL) {

		if (b->b_instr)

			PyObject_Free((void *)b->b_instr);

		next = b->b_list;

		PyObject_Free((void *)b);

		b = next;

	}

	Py_CLEAR(u->u_ste);

	Py_CLEAR(u->u_name);

	Py_CLEAR(u->u_consts);

	Py_CLEAR(u->u_names);

	Py_CLEAR(u->u_varnames);

	Py_CLEAR(u->u_freevars);

	Py_CLEAR(u->u_cellvars);

	Py_CLEAR(u->u_private);

	PyObject_Free(u);

}

static int

compiler_enter_scope(struct compiler *c, identifier name, void *key,

		     int lineno)

{

	struct compiler_unit *u;

	u = (struct compiler_unit *)PyObject_Malloc(sizeof(

						struct compiler_unit));

	if (!u) {

		PyErr_NoMemory();

		return 0;

	}

	memset(u, 0, sizeof(struct compiler_unit));

	u->u_argcount = 0;

	u->u_ste = PySymtable_Lookup(c->c_st, key);

	if (!u->u_ste) {

		compiler_unit_free(u);

		return 0;

	}

	Py_INCREF(name);

	u->u_name = name;

	u->u_varnames = list2dict(u->u_ste->ste_varnames);

	u->u_cellvars = dictbytype(u->u_ste->ste_symbols, CELL, 0, 0);

	if (!u->u_varnames || !u->u_cellvars) {

		compiler_unit_free(u);

		return 0;

	}

	u->u_freevars = dictbytype(u->u_ste->ste_symbols, FREE, DEF_FREE_CLASS,

				   PyDict_Size(u->u_cellvars));

	if (!u->u_freevars) {

		compiler_unit_free(u);

		return 0;

	}

	u->u_blocks = NULL;

	u->u_tmpname = 0;

	u->u_nfblocks = 0;

	u->u_firstlineno = lineno;

	u->u_lineno = 0;

	u->u_lineno_set = false;

	u->u_consts = PyDict_New();

	if (!u->u_consts) {

		compiler_unit_free(u);

		return 0;

	}

	u->u_names = PyDict_New();

	if (!u->u_names) {

		compiler_unit_free(u);

		return 0;

	}

	u->u_private = NULL;

	/* Push the old compiler_unit on the stack. */

	if (c->u) {

		PyObject *wrapper = PyCObject_FromVoidPtr(c->u, NULL);

		if (!wrapper || PyList_Append(c->c_stack, wrapper) < 0) {

			Py_XDECREF(wrapper);

			compiler_unit_free(u);

			return 0;

		}

		Py_DECREF(wrapper);

		u->u_private = c->u->u_private;

		Py_XINCREF(u->u_private);

	}

	c->u = u;

	c->c_nestlevel++;

	if (compiler_use_new_block(c) == NULL)

		return 0;

	return 1;

}

static void

compiler_exit_scope(struct compiler *c)

{

	int n;

	PyObject *wrapper;

	c->c_nestlevel--;

	compiler_unit_free(c->u);

	/* Restore c->u to the parent unit. */

	n = PyList_GET_SIZE(c->c_stack) - 1;

	if (n >= 0) {

		wrapper = PyList_GET_ITEM(c->c_stack, n);

		c->u = (struct compiler_unit *)PyCObject_AsVoidPtr(wrapper);

		assert(c->u);

		/* we are deleting from a list so this really shouldn't fail */