count()

global.c

@@ -165,7 +165,22 @@ static int global_deref_type_infer(const struct func *func, struct node *n)
 
 /* :map */
 
-static int global_map_ir_pre_key(struct node *n, struct prog *prog)
+static int map_ir_update(struct node *n, struct prog *prog)
+{
+	struct node *map = n->expr.args;
+
+	ir_emit_ldmap(prog->ir, BPF_REG_1, map->sym);
+	ir_emit_insn(prog->ir, MOV, BPF_REG_2, BPF_REG_BP);
+	ir_emit_insn(prog->ir, ALU_IMM(BPF_ADD, map->sym->irs.stack), BPF_REG_2, 0);
+	ir_emit_insn(prog->ir, MOV, BPF_REG_3, BPF_REG_BP);
+	ir_emit_insn(prog->ir, ALU_IMM(BPF_ADD, n->sym->irs.stack), BPF_REG_3, 0);
+	ir_emit_insn(prog->ir, MOV_IMM(0), BPF_REG_4, 0);
+	ir_emit_insn(prog->ir, CALL(BPF_FUNC_map_update_elem), 0, 0);
+	/* TODO: if (r0) exit(r0); */
+	return 0;
+}
+
+static int map_ir_pre_key(struct node *n, struct prog *prog)
 {
 	struct node *map = n->expr.args, *arg;
 	struct type *ktype = type_base(map->sym->type->map.ktype);
@@ -198,8 +213,8 @@ static int global_map_ir_pre_key(struct node *n, struct prog *prog)
 	return 0;
 }
 
-static int global_map_ir_pre(const struct func *func, struct node *n,
-			     struct prog *prog)
+static int map_ir_pre(const struct func *func, struct node *n,
+		      struct prog *prog)
 {
 	struct irstate *kirs;
 	struct node *map = n->expr.args;
@@ -210,7 +225,7 @@ static int global_map_ir_pre(const struct func *func, struct node *n,
 
 
 	if (ktype->ttype == T_STRUCT)
-		return global_map_ir_pre_key(n, prog);
+		return map_ir_pre_key(n, prog);
 
 	kirs = &map->next->sym->irs;
 	if (!kirs->loc) {
@@ -220,8 +235,8 @@ static int global_map_ir_pre(const struct func *func, struct node *n,
 	return 0;
 }
 
-static int global_map_ir_post(const struct func *func, struct node *n,
-			      struct prog *prog)
+static int map_ir_post(const struct func *func, struct node *n,
+		       struct prog *prog)
 {
 	struct node *map = n->expr.args, *arg;
 	struct type *ktype = type_base(map->sym->type->map.ktype);
@@ -231,10 +246,16 @@ static int global_map_ir_post(const struct func *func, struct node *n,
 	int16_t lmiss, lhit;
 
 	arg = map->next;
-	tfields_foreach(f, ktype->sou.fields) {
-		offset = type_offsetof(ktype, f->name);
-		ir_emit_sym_to_stack(prog->ir, stack + offset, arg->sym);
-		arg = arg->next;
+
+	if (ktype->ttype == T_STRUCT) {
+		tfields_foreach(f, ktype->sou.fields) {
+			offset = type_offsetof(ktype, f->name);
+			ir_emit_sym_to_stack(prog->ir, stack + offset, arg->sym);
+			arg = arg->next;
+		}
+	} else {
+		ir_emit_sym_to_stack(prog->ir, stack, arg->sym);
+		assert(!arg->next);
 	}
 
 	n->sym->irs.hint.stack = 1;
@@ -264,7 +285,7 @@ static int global_map_ir_post(const struct func *func, struct node *n,
 	return 0;
 }
 
-static struct type *global_map_ktype(struct node *n)
+static struct type *map_key_type(struct node *n)
 {
 	struct node *map, *key;
 	struct type *ktype;
@@ -296,22 +317,44 @@ static struct type *global_map_ktype(struct node *n)
 	return ktype;
 }
 
-static int global_map_type_infer(const struct func *func, struct node *n)
+static int map_type_validate(struct node *n)
 {
-	struct node *map = n->expr.args;
+	/* TODO */
+	return 0;
+}
 
-	if (!map->sym->type)
+static int map_type_infer(const struct func *func, struct node *n)
+{
+	struct node *map, *key;
+	struct type *ktype;
+
+	map = n->expr.args;
+	if (!map->sym)
 		return 0;
 
-	/* TODO validate key against known type */
+	if (map->sym->type) {
+		if (!n->sym->type)
+			/* given `m[key]` where m's type is known,
+			 * infer that the expression's type is equal
+			 * to m's value type. */
+			n->sym->type = map->sym->type->map.vtype;
+
+		return map_type_validate(n);
+	}
 
-	/* given `m[key]` where m's type is known, infer that the
-	 * expression's type is equal to m's value type. */
-	n->sym->type = map->sym->type->map.vtype;
+	if (!n->sym->type)
+		return 0;
+
+	for (key = map->next; key; key = key->next) {
+		if (type_sizeof(key->sym->type) < 0)
+			return 0;
+	}
+
+	map->sym->type = type_map_of(map_key_type(n), n->sym->type);
 	return 0;
 }
 
-static int global_map_static_validate(const struct func *func, struct node *n)
+static int map_static_validate(const struct func *func, struct node *n)
 {
 	if (n->expr.args->ntype != N_IDENT) {
 		_e("%#N: can't lookup a key in %N, which is not a map.\n",
@@ -347,7 +390,7 @@ static int global_assign_ir_pre(const struct func *func, struct node *n,
 static int global_assign_ir_post(const struct func *func, struct node *n,
 				 struct prog *prog)
 {
-	struct node *lval, *rval, *map;
+	struct node *lval, *rval;
 
 	lval = n->expr.args;
 	rval = lval->next;
@@ -356,38 +399,13 @@ static int global_assign_ir_post(const struct func *func, struct node *n,
 	if (!node_is(lval, "{}"))
 		return 0;
 
-	map = lval->expr.args;
-
-	ir_emit_ldmap(prog->ir, BPF_REG_1, map->sym);
-	ir_emit_insn(prog->ir, MOV, BPF_REG_2, BPF_REG_BP);
-	ir_emit_insn(prog->ir, ALU_IMM(BPF_ADD, map->sym->irs.stack), BPF_REG_2, 0);
-	ir_emit_insn(prog->ir, MOV, BPF_REG_3, BPF_REG_BP);
-	ir_emit_insn(prog->ir, ALU_IMM(BPF_ADD, lval->sym->irs.stack), BPF_REG_3, 0);
-	ir_emit_insn(prog->ir, MOV_IMM(0), BPF_REG_4, 0);
-	ir_emit_insn(prog->ir, CALL(BPF_FUNC_map_update_elem), 0, 0);
-	/* TODO: if (r0) exit(r0); */
-	return 0;
-}
-
-static int global_assign_type_infer_map(struct node *n)
-{
-	struct node *map, *key;
-	struct type *ktype;
-
-	map = n->expr.args;
-
-	for (key = map->next; key; key = key->next) {
-		if (type_sizeof(key->sym->type) < 0)
-			return 0;
-	}
-
-	map->sym->type = type_map_of(global_map_ktype(n), n->sym->type);
-	return 0;
+	return map_ir_update(lval, prog);
 }
 
 static int global_assign_type_infer(const struct func *func, struct node *n)
 {
 	struct node *lval, *rval;
+	int err;
 
 	if (n->sym->type)
 		return 0;
@@ -406,10 +424,12 @@ static int global_assign_type_infer(const struct func *func, struct node *n)
 		/* TODO do we need assignment expressions? */
 		n->sym->type = &t_void;
 		
-		if (node_is(lval, "{}"))
-			return global_assign_type_infer_map(lval);
+		if (!node_is(lval, "{}"))
+			return 0;
 
-		return 0;
+		err = map_type_infer(lval->sym->func, lval);
+		if (err)
+			return err;
 	}
 
 	if (type_compatible(lval->sym->type, rval->sym->type))
@@ -458,7 +478,68 @@ static int global_binop_type_infer(const struct func *func, struct node *n)
 	return 0;
 }
 
-/* :binop */
+
+/* count() */
+
+struct type t_count = {
+	.ttype = T_SCALAR,
+
+	.scalar = {
+		.name = ":count",
+		.size = sizeof(unsigned long),
+		.is_signed = 0,
+	},
+
+	.aggregation = 1,
+};
+
+static int count_ir_post(const struct func *func, struct node *n,
+			 struct prog *prog)
+{
+	struct node *mapop = n->expr.args;
+
+	ir_emit_sym_to_reg(prog->ir, BPF_REG_0, mapop->sym);
+	ir_emit_insn(prog->ir, ALU_IMM(BPF_ADD, 1), BPF_REG_0, 0);
+	ir_emit_reg_to_sym(prog->ir, mapop->sym, BPF_REG_0);
+
+	return map_ir_update(mapop, prog);
+}
+
+static int count_type_infer(const struct func *func, struct node *n)
+{
+	struct node *mapop;
+	struct type *t;
+
+	mapop = n->expr.args;
+
+	if (n->sym->type)
+		return 0;
+
+	n->sym->type = &t_void;
+
+	if (!mapop->sym->type) {
+		mapop->sym->type = &t_count;
+		return map_type_infer(mapop->sym->func, mapop);
+	}
+
+	return 0;
+}
+
+static int count_static_validate(const struct func *func, struct node *n)
+{
+	struct node *mapop;
+
+	mapop = n->expr.args;
+
+	if (node_is(mapop, "{}"))
+		return 0;
+
+	_e("%#N: aggregation target must be a map, %N is not.\n", n, mapop);
+	return -EINVAL;
+}
+
+
+/* quantize */
 
 static int global_quantize_type_infer(const struct func *func, struct node *n)
 {
@@ -640,13 +721,20 @@ static const struct func global_funcs[] = {
 	{
 		.name = "{}",
 		/* .type = t_map_func, */
-		.type_infer = global_map_type_infer,
-		.static_validate = global_map_static_validate,
+		.type_infer = map_type_infer,
+		.static_validate = map_static_validate,
 
-		.ir_pre  = global_map_ir_pre,
-		.ir_post = global_map_ir_post,
+		.ir_pre  = map_ir_pre,
+		.ir_post = map_ir_post,
 	},
+	{
+		.name = "count",
+		.type = &t_1arg_func,
+		.static_validate = count_static_validate,
+		.type_infer = count_type_infer,
 
+		.ir_post = count_ir_post,
+	},
 	{
 		.name = "pid",
 		.type = &t_pid_func,

ir.c

@@ -180,9 +180,12 @@ static void __insn_dump(const struct bpf_insn insn, uint16_t dst, uint16_t src,
 
 	if (BPF_CLASS(insn.code) == BPF_LDX || BPF_CLASS(insn.code) == BPF_STX)
 		goto reg_src;
+	else if (BPF_CLASS(insn.code) == BPF_ST)
+		goto imm_src;
 
 	switch (BPF_SRC(insn.code)) {
 	case BPF_K:
+	imm_src:
 		fprintf(fp, "#%s0x%x", insn.imm < 0 ? "-" : "",
 			insn.imm < 0 ? -insn.imm : insn.imm);
 		break;
@@ -530,3 +533,10 @@ struct ir *ir_new(void)
 	ir->next_label = -1;
 	return ir;
 }
+
+
+
+int ir_generate_bpf(struct ir *ir)
+{
+	return 0;
+}

ir.h

@@ -153,4 +153,6 @@ void ir_emit_bzero(struct ir *ir, ssize_t offset, size_t size);
 
 struct ir *ir_new(void);
 
+int ir_generate_bpf(struct ir *ir);
+
 #endif	/* _PLY_IR_H */

ply.c

@@ -82,14 +82,21 @@ struct ctx *ctx_get(void)
 			  /* 	    node_ident("us"), */
 			  /* 	    node_expr("pid", NULL), */
 			  /* 	    NULL), */
-			  node_expr("=",
-				    node_expr("{}",
-					      node_ident("t"),
-					      node_expr("pid", NULL),
-					      node_num("0x10000000"),
+			  /* node_expr("=", */
+			  /* 	    node_expr("{}", */
+			  /* 		      node_ident("t"), */
+			  /* 		      node_expr("pid", NULL), */
+			  /* 		      node_num("0x10000000"), */
+			  /* 		      NULL), */
+			  /* 	    node_expr("time", NULL), */
+			  /* 	    NULL), */
+			  node_expr("count",
+			  	    node_expr("{}",
+			  		      node_ident("reads"),
+			  		      node_expr("pid", NULL),
+					      node_ident("arg0"),
 					      NULL),
-				    node_expr("time", NULL),
-				    NULL),
+			  	    NULL),
 			  /* node_expr("=", */
 			  /* 	    node_expr("{}", */
 			  /* 		      node_ident("reads"), */
@@ -213,6 +220,20 @@ int pass_ir_post(struct node *n, void *_prog)
 	return 0;
 }
 
+int run_walk(struct pass *pass, struct ctx *ctx)
+{
+	struct prog **prog;
+	int err;
+
+	for (prog = ctx->progs; *prog; prog++) {
+		err = node_walk((*prog)->ast, pass->pre, pass->post, *prog);
+		if (err)
+			return err;
+	}
+
+	return 0;
+}
+
 int run_ir(struct pass *pass, struct ctx *ctx)
 {
 	struct prog **progp;
@@ -241,13 +262,15 @@ int run_ir(struct pass *pass, struct ctx *ctx)
 	return 0;
 }
 
-int run_walk(struct pass *pass, struct ctx *ctx)
+int run_bpf(struct pass *pass, struct ctx *ctx)
 {
-	struct prog **prog;
+	struct prog **progp;
 	int err;
 
-	for (prog = ctx->progs; *prog; prog++) {
-		err = node_walk((*prog)->ast, pass->pre, pass->post, *prog);
+	for (progp = ctx->progs; *progp; progp++) {
+		struct prog *prog = *progp;
+
+		err = ir_generate_bpf(prog->ir);
 		if (err)
 			return err;
 	}

type.c

@@ -31,8 +31,7 @@ static void type_dump_func(struct type *t, const char *name, FILE *fp)
 	struct tfield *arg;
 
 	type_dump(t->func.type, NULL, fp);
-	fprintf(fp, " %s(*\e[1m%s\e[0m)(",
-		t->func.aggregation ? "@" : "", name ? : "");
+	fprintf(fp, " (*\e[1m%s\e[0m)(", name ? : "");
 
 	if (!t->func.args) {
 		__faint(fp, t->func.vargs ? "..." : "void");
@@ -60,6 +59,9 @@ void type_dump(struct type *t, const char *name, FILE *fp)
 		goto print_name;
 	}
 
+	if (t->aggregation)
+		fputc('@', fp);
+
 	switch (t->ttype){
 	case T_VOID:
 		__faint(fp, "void");

type.h

@@ -52,7 +52,6 @@ struct tfunc {
 	struct tfield *args;
 
 	int vargs:1;
-	int aggregation:1;
 };
 
 enum ttype {
@@ -78,6 +77,8 @@ struct type {
 		struct tstruct sou;
 		struct tfunc func;
 	};
+
+	int aggregation:1;
 };
 
 int type_equal     (struct type *a, struct type *b);