#ifndef _PLY_IR_H
#define _PLY_IR_H

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#include <linux/bpf.h>

#define INSN(_code, _dst, _src, _off, _imm)	\
	((struct bpf_insn) {			\
		.code  = _code,			\
		.dst_reg = _dst,		\
		.src_reg = _src,		\
		.off   = _off,			\
		.imm   = _imm			\
	})

#define MOV(_dst, _src)     INSN(BPF_ALU64 | BPF_MOV | BPF_X, _dst, _src, 0, 0)
#define MOV_IMM(_dst, _imm) INSN(BPF_ALU64 | BPF_MOV | BPF_K, _dst, 0, 0, _imm)

#define EXIT INSN(BPF_JMP | BPF_EXIT, 0, 0, 0, 0)
#define CALL(_imm) INSN(BPF_JMP | BPF_CALL, 0, 0, 0, _imm)
#define JMP(_op, _dst, _src, _off)     INSN(BPF_JMP | BPF_OP((_op)) | BPF_X, _dst, _src, _off, 0)
#define JMP_IMM(_op, _dst, _imm, _off) INSN(BPF_JMP | BPF_OP((_op)) | BPF_K, _dst, 0, _off, _imm)

#define ALU(_op, _dst, _src)     INSN(BPF_ALU64 | BPF_OP((_op)) | BPF_X, _dst, _src, 0, 0)
#define ALU_IMM(_op, _dst, _imm) INSN(BPF_ALU64 | BPF_OP((_op)) | BPF_K, _dst, 0, 0, _imm)

#define STW_IMM(_dst, _off, _imm) INSN(BPF_ST  | BPF_SIZE(BPF_W)  | BPF_MEM, _dst, 0, _off, _imm)
#define STXB(_dst, _off, _src)   INSN(BPF_STX | BPF_SIZE(BPF_B) | BPF_MEM, _dst, _src, _off, 0)
#define STXH(_dst, _off, _src)   INSN(BPF_STX | BPF_SIZE(BPF_H) | BPF_MEM, _dst, _src, _off, 0)
#define STXW(_dst, _off, _src)   INSN(BPF_STX | BPF_SIZE(BPF_W) | BPF_MEM, _dst, _src, _off, 0)
#define STXDW(_dst, _off, _src)   INSN(BPF_STX | BPF_SIZE(BPF_DW) | BPF_MEM, _dst, _src, _off, 0)

#define LDXB(_dst, _off, _src)  INSN(BPF_LDX | BPF_SIZE(BPF_B)  | BPF_MEM, _dst, _src, _off, 0)
#define LDXH(_dst, _off, _src)  INSN(BPF_LDX | BPF_SIZE(BPF_H)  | BPF_MEM, _dst, _src, _off, 0)
#define LDXW(_dst, _off, _src)  INSN(BPF_LDX | BPF_SIZE(BPF_W)  | BPF_MEM, _dst, _src, _off, 0)
#define LDXDW(_dst, _off, _src) INSN(BPF_LDX | BPF_SIZE(BPF_DW) | BPF_MEM, _dst, _src, _off, 0)

struct sym;

enum vitype {
	VI_INSN,
	VI_LDMAP,
	VI_LABEL,
	VI_REG_GET,
	VI_REG_PUT,
};

struct vinsn {
	enum vitype vitype;

	union {
		struct {
			struct bpf_insn bpf;
			uint16_t dst;
			uint16_t src;
		} insn;

		int16_t label;
		uint16_t reg;

		struct {
			uint16_t reg;
			struct sym *sym;
		} map;
	};
};

struct ir {
	struct vinsn *vi;
	size_t len;

	int16_t next_label;
	uint16_t next_reg;
};

void insn_dump(struct bpf_insn insn, FILE *fp);
void vinsn_dump(struct vinsn *vi, FILE *fp);
void ir_dump(struct ir *ir, FILE *fp);

int16_t  ir_alloc_label   (struct ir *ir);
uint16_t ir_alloc_register(struct ir *ir);

void ir_emit_insn   (struct ir *ir, struct bpf_insn bpf, uint16_t dst, uint16_t src);
void ir_emit_ldmap  (struct ir *ir, uint16_t dst, struct sym *map);
void ir_emit_label  (struct ir *ir, int16_t label);
void ir_emit_reg_get(struct ir *ir, uint16_t reg);
void ir_emit_reg_put(struct ir *ir, uint16_t reg);

static inline void ir_emit_call(struct ir *ir, enum bpf_func_id func)
{
	ir_emit_reg_get(ir, BPF_REG_0);
	ir_emit_reg_get(ir, BPF_REG_1);
	ir_emit_reg_get(ir, BPF_REG_2);
	ir_emit_reg_get(ir, BPF_REG_3);
	ir_emit_reg_get(ir, BPF_REG_4);
	ir_emit_reg_get(ir, BPF_REG_5);
	ir_emit_insn(ir, CALL(func), 0, 0);
	ir_emit_reg_put(ir, BPF_REG_5);
	ir_emit_reg_put(ir, BPF_REG_4);
	ir_emit_reg_put(ir, BPF_REG_3);
	ir_emit_reg_put(ir, BPF_REG_2);
	ir_emit_reg_put(ir, BPF_REG_1);
	ir_emit_reg_put(ir, BPF_REG_0);
}

struct ir *ir_new(void);

enum irloc {
	LOC_IMM   = (1 << 0),
	LOC_REG   = (1 << 1),
	LOC_STACK = (1 << 2),
};

struct irstate {
	int loc;

	uint16_t reg;
	int16_t stack;
};

static inline void irs_alloc_reg(struct irstate *irs, struct ir *ir)
{
	irs->loc = LOC_REG;
	irs->reg = ir_alloc_register(ir);
	ir_emit_reg_get(ir, irs->reg);
}

#endif	/* _PLY_IR_H */