diff --git a/lang/ast/ctx.c b/lang/ast/ctx.c
index df9d1d8..0b479ad 100644
--- a/lang/ast/ctx.c
+++ b/lang/ast/ctx.c
@@ -18,8 +18,7 @@ enum ivy_status ivy_parser_create(struct ivy_parser **parser)
 
 	memset(out, 0x0, sizeof *out);
 
-	parser_push_state(
-		out, IVY_AST_UNIT, struct ivy_ast_unit_node, struct parser_state);
+	parser_push_state(		out, IVY_AST_UNIT, struct ivy_ast_unit_node);
 
 	*parser = out;
 	return IVY_OK;
@@ -67,14 +66,19 @@ struct parser_state *parser_get_state_generic(struct ivy_parser *parser)
 
 struct parser_state *parser_push_state_generic(
 	struct ivy_parser *parser, enum ivy_ast_node_type type,
-	size_t node_size, size_t state_size)
+	size_t node_size)
 {
-	struct parser_state *state = malloc(state_size);
+	const struct ast_node_type *node_type = get_ast_node_type(type);
+	if (!node_type) {
+		return NULL;
+	}
+
+	struct parser_state *state = malloc(node_type->n_state_size);
 	if (!state) {
 		return NULL;
 	}
 
-	memset(state, 0x0, state_size);
+	memset(state, 0x0, node_type->n_state_size);
 
 	b_queue_entry *current_state_entry = b_queue_last(&parser->p_state);
 	if (current_state_entry) {
@@ -85,8 +89,7 @@ struct parser_state *parser_push_state_generic(
 
 	state->s_node = ast_node_create_with_size(type, node_size);
 
-	const struct ast_node_type *node_type = get_ast_node_type(type);
-	if (node_type && node_type->n_init_state) {
+	if (node_type->n_init_state) {
 		node_type->n_init_state(state);
 	}
 
diff --git a/lang/ast/ctx.h b/lang/ast/ctx.h
index ca099b7..87d21ff 100644
--- a/lang/ast/ctx.h
+++ b/lang/ast/ctx.h
@@ -5,9 +5,9 @@
 #include <ivy/lang/ast.h>
 #include <ivy/status.h>
 
-#define parser_push_state(parser, node_id, node_type, state_type)              \
-	((state_type *)parser_push_state_generic(                              \
-		parser, node_id, sizeof(node_type), sizeof(state_type)))
+#define parser_push_state(parser, node_id, node_type)              \
+	(parser_push_state_generic(                              \
+		parser, node_id, sizeof(node_type)))
 #define parser_get_state(parser, state_type)                                   \
 	((state_type *)parser_get_state_generic(parser))
 
@@ -30,7 +30,7 @@ enum pop_state_flags {
 
 extern struct parser_state *parser_push_state_generic(
 	struct ivy_parser *parser, enum ivy_ast_node_type node_type,
-	size_t node_size, size_t state_size);
+	size_t node_size);
 extern void parser_pop_state(struct ivy_parser *parser, enum pop_state_flags flags);
 extern struct parser_state *parser_get_state_generic(struct ivy_parser *parser);
 
diff --git a/lang/ast/node.h b/lang/ast/node.h
index fcc4dcf..143541c 100644
--- a/lang/ast/node.h
+++ b/lang/ast/node.h
@@ -17,6 +17,7 @@ struct ast_node_type {
 		struct ivy_ast_node *, struct ivy_ast_node *);
 	void (*n_print)(struct ivy_ast_node *);
 	void (*n_init_state)(struct parser_state *);
+	size_t n_state_size;
 
 	token_parse_function n_token_parsers[IVY_TOK_TYPE_COUNT];
 	token_parse_function n_keyword_parsers[IVY_KW_TYPE_COUNT];
diff --git a/lang/ast/unit-package.c b/lang/ast/unit-package.c
index 000b9c2..fd95c68 100644
--- a/lang/ast/unit-package.c
+++ b/lang/ast/unit-package.c
@@ -62,6 +62,7 @@ struct ast_node_type unit_package_node_ops = {
 	.n_add_child = add_child,
         .n_print = print,
 	.n_init_state = init_state,
+	.n_state_size = sizeof(struct unit_package_parser_state),
 	.n_symbol_parsers = {
 		[IVY_SYM_DOT] = parse_dot,
 	},
diff --git a/lang/ast/unit.c b/lang/ast/unit.c
index 5e75b41..0807b91 100644
--- a/lang/ast/unit.c
+++ b/lang/ast/unit.c
@@ -6,8 +6,7 @@ static enum ivy_status parse_package_keyword(
 	struct ivy_parser *ctx, struct ivy_token *tok)
 {
 	parser_push_state(
-		ctx, IVY_AST_UNIT_PACKAGE, struct ivy_ast_unit_package_node,
-		struct unit_package_parser_state);
+		ctx, IVY_AST_UNIT_PACKAGE, struct ivy_ast_unit_package_node);
 	return IVY_OK;
 }
 
@@ -24,6 +23,7 @@ static void print(struct ivy_ast_node *node)
 struct ast_node_type unit_node_ops = {
 	.n_add_child = add_child,
         .n_print = print,
+	.n_state_size = sizeof(struct parser_state),
         .n_keyword_parsers = {
 		[IVY_KW_PACKAGE] = parse_package_keyword,
 	},