Refactoring

src/dune

@@ -1,6 +1,6 @@
 (library
  (name stellogen)
- (libraries base menhirLib)
+ (libraries base stdio menhirLib)
  (preprocess
   (pps sedlex.ppx ppx_deriving.show ppx_deriving.ord ppx_deriving.eq)))
 

src/expr.ml

@@ -3,6 +3,8 @@ open Lsc_ast
 open Sgen_ast
 open Expr_err
 
+exception MacroError of expr_err * source_location option
+
 let ( let* ) x f = Result.bind x ~f
 
 type ident = string
@@ -138,6 +140,141 @@ let rec replace_id (var_from : ident) replacement (expr : expr loc) : expr loc =
     ; loc = expr.loc
     }
 
+(* ---------------------------------------
+   Macro Expansion Helpers
+   --------------------------------------- *)
+
+(* Check if a pattern's formal params indicate variadic (ends with Var ...) *)
+let is_variadic_pattern (formal_params : string list) : bool =
+  match List.rev formal_params with "..." :: _ :: _ -> true | _ -> false
+
+(* Calculate minimum arguments for a pattern *)
+let min_args_for_pattern (formal_params : string list) : int =
+  if is_variadic_pattern formal_params then
+    match List.rev formal_params with
+    | "..." :: _var :: rest -> List.length rest
+    | _ -> List.length formal_params
+  else List.length formal_params
+
+(* Check if a pattern matches the given argument count *)
+let pattern_matches_args (formal_params : string list) (arg_count : int) : bool
+    =
+  let min_args = min_args_for_pattern formal_params in
+  if is_variadic_pattern formal_params then arg_count >= min_args
+  else arg_count = min_args
+
+(* Split formal params into fixed params and optional variadic param name *)
+let split_variadic_params (formal_params : string list) :
+  string list * string option =
+  if is_variadic_pattern formal_params then
+    match List.rev formal_params with
+    | "..." :: var :: rest -> (List.rev rest, Some var)
+    | _ -> (formal_params, None)
+  else (formal_params, None)
+
+(* Find the best matching pattern for a macro name.
+   Prefers exact (non-variadic) matches over variadic ones. *)
+let find_matching_pattern
+  (all_patterns : (string * (string list * expr loc list)) list)
+  (arg_count : int) : (string list * expr loc list) option =
+  (* First try exact (non-variadic) matches *)
+  match
+    List.find all_patterns ~f:(fun (_, (params, _)) ->
+      pattern_matches_args params arg_count && not (is_variadic_pattern params) )
+  with
+  | Some (_, pattern) -> Some pattern
+  | None ->
+    (* Then try variadic matches *)
+    List.find_map all_patterns ~f:(fun ((_, (params, _)) as pattern) ->
+      if pattern_matches_args params arg_count then Some (snd pattern) else None )
+
+(* Apply substitution to an expression, handling variadic splicing.
+   - subst_pairs: fixed param -> arg mappings
+   - variadic_param: optional name of variadic param
+   - rest_args: arguments captured by variadic param *)
+let rec apply_variadic_substitution (subst_pairs : (string * expr loc) list)
+  (variadic_param : string option) (rest_args : expr loc list) (e : expr loc) :
+  expr loc =
+  match (e.content, variadic_param) with
+  | List subexprs, Some vp ->
+    (* Scan for Var ... pattern and splice *)
+    let rec splice_list exprs =
+      match exprs with
+      | { content = Var v; _ } :: { content = Symbol "..."; _ } :: rest
+        when String.equal v vp ->
+        (* Found the pattern - splice in rest_args at this position *)
+        rest_args @ splice_list rest
+      | e :: rest ->
+        (* Recursively substitute in this element, then continue *)
+        apply_variadic_substitution subst_pairs variadic_param rest_args e
+        :: splice_list rest
+      | [] -> []
+    in
+    { content = List (splice_list subexprs); loc = e.loc }
+  | List subexprs, None ->
+    (* No variadic param - just recurse *)
+    { content =
+        List
+          (List.map subexprs
+             ~f:
+               (apply_variadic_substitution subst_pairs variadic_param rest_args) )
+    ; loc = e.loc
+    }
+  | Var v, Some vp when String.equal v vp ->
+    (* Variadic param used alone without ... - wrap rest_args in a list *)
+    { content = List rest_args; loc = e.loc }
+  | Var v, _ -> (
+    match List.Assoc.find subst_pairs v ~equal:String.equal with
+    | Some replacement -> { replacement with loc = e.loc }
+    | None -> e )
+  | _, _ -> e
+
+(* Expand a single matched macro: substitute params and recursively expand *)
+let expand_matched_macro
+  (macro_env : (string * (string list * expr loc list)) list)
+  (expand_fn :
+    (string * (string list * expr loc list)) list -> expr loc -> expr loc list
+    ) (formal_params : string list) (body : expr loc list)
+  (call_args : expr loc list) (call_loc : source_location option) :
+  expr loc list =
+  let fixed_params, variadic_param = split_variadic_params formal_params in
+  let min_args = List.length fixed_params in
+  (* First, recursively expand macros in the arguments *)
+  let expanded_args =
+    List.map call_args ~f:(fun arg ->
+      match expand_fn macro_env arg with
+      | [ single ] -> single
+      | multiple -> { content = List multiple; loc = arg.loc } )
+  in
+  (* Split into fixed and variadic args *)
+  let fixed_args, rest_args = List.split_n expanded_args min_args in
+  (* Build substitution environment *)
+  let subst_pairs = List.zip_exn fixed_params fixed_args in
+  (* Apply substitution to body *)
+  let substituted =
+    List.map body
+      ~f:(apply_variadic_substitution subst_pairs variadic_param rest_args)
+  in
+  (* Recursively expand macros in the substituted body *)
+  let expanded = List.concat_map substituted ~f:(expand_fn macro_env) in
+  (* Attach the call site location to the expanded expressions *)
+  List.map expanded ~f:(fun e -> { e with loc = call_loc })
+
+(* Expand macros in a list of arguments, coalescing results *)
+let expand_args_in_list
+  (macro_env : (string * (string list * expr loc list)) list)
+  (expand_fn :
+    (string * (string list * expr loc list)) list -> expr loc -> expr loc list
+    ) (args : expr loc list) : expr loc list =
+  List.map args ~f:(fun arg ->
+    match expand_fn macro_env arg with
+    | [ single ] -> single
+    | multiple -> { content = List multiple; loc = arg.loc } )
+
+(* ---------------------------------------
+   Main Macro Expansion
+   --------------------------------------- *)
+
 (* Recursively expand macros in an expression *)
 let rec expand_macros_in_expr
   (macro_env : (string * (string list * expr loc list)) list) (expr : expr loc)
@@ -155,112 +292,15 @@ let rec expand_macros_in_expr
     let all_patterns =
       List.filter macro_env ~f:(fun (name, _) -> String.equal name macro_name)
     in
-
-    (* Helper to check if a pattern matches the argument count *)
-    let pattern_matches (formal_params, _body) =
-      let is_variadic =
-        match List.rev formal_params with "..." :: _ :: _ -> true | _ -> false
-      in
-      let min_args =
-        if is_variadic then
-          match List.rev formal_params with
-          | "..." :: _var :: rest -> List.length rest
-          | _ -> List.length formal_params
-        else List.length formal_params
-      in
-      let arg_count = List.length call_args in
-      if is_variadic then arg_count >= min_args else arg_count = min_args
-    in
-
-    (* Find the best matching pattern: prefer exact matches (non-variadic) over variadic *)
-    let matching_pattern =
-      (* First try exact matches *)
-      match
-        List.find all_patterns ~f:(fun (_, pattern) ->
-          pattern_matches pattern
-          && not
-               ( match List.rev (fst pattern) with
-               | "..." :: _ :: _ -> true
-               | _ -> false ) )
-      with
-      | Some (_, pattern) -> Some pattern
-      | None ->
-        (* Then try variadic matches *)
-        List.find_map all_patterns ~f:(fun (_, pattern) ->
-          if pattern_matches pattern then Some pattern else None )
-    in
-
-    match matching_pattern with
+    let arg_count = List.length call_args in
+    match find_matching_pattern all_patterns arg_count with
     | Some (formal_params, body) ->
-      (* Check if this is a variadic macro (ends with ...) *)
-      let is_variadic =
-        match List.rev formal_params with "..." :: _ :: _ -> true | _ -> false
-      in
-      let fixed_params, variadic_param =
-        if is_variadic then
-          match List.rev formal_params with
-          | "..." :: var :: rest -> (List.rev rest, Some var)
-          | _ -> (formal_params, None)
-        else (formal_params, None)
-      in
-      let min_args = List.length fixed_params in
-      (* First, recursively expand macros in the arguments *)
-      let expanded_args =
-        List.map call_args ~f:(fun arg ->
-          match expand_macros_in_expr macro_env arg with
-          | [ single ] -> single
-          | multiple -> { content = List multiple; loc = arg.loc } )
-      in
-      (* Split into fixed and variadic args *)
-      let fixed_args, rest_args = List.split_n expanded_args min_args in
-      (* Build substitution environment *)
-      let subst_pairs = List.zip_exn fixed_params fixed_args in
-      (* Helper to substitute with variadic support *)
-      let rec apply_substitution_var e =
-        match (e.content, variadic_param) with
-        | List subexprs, Some vp ->
-          (* Scan for Var ... pattern and splice *)
-          let rec splice_list exprs =
-            match exprs with
-            | { content = Var v; _ } :: { content = Symbol "..."; _ } :: rest
-              when String.equal v vp ->
-              (* Found the pattern - splice in rest_args at this position *)
-              rest_args @ splice_list rest
-            | e :: rest ->
-              (* Recursively substitute in this element, then continue *)
-              apply_substitution_var e :: splice_list rest
-            | [] -> []
-          in
-          { content = List (splice_list subexprs); loc = e.loc }
-        | List subexprs, None ->
-          (* No variadic param - just recurse *)
-          { content = List (List.map subexprs ~f:apply_substitution_var)
-          ; loc = e.loc
-          }
-        | Var v, Some vp when String.equal v vp ->
-          (* Variadic param used alone without ... - this is an error in most cases,
-             but we handle it by wrapping rest_args in a list *)
-          { content = List rest_args; loc = e.loc }
-        | Var v, _ -> (
-          match List.Assoc.find subst_pairs v ~equal:String.equal with
-          | Some replacement -> { replacement with loc = e.loc }
-          | None -> e )
-        | _, _ -> e
-      in
-      let substituted = List.map body ~f:apply_substitution_var in
-      (* Recursively expand macros in the substituted body *)
-      let expanded =
-        List.concat_map substituted ~f:(expand_macros_in_expr macro_env)
-      in
-      (* Attach the call site location to the expanded expressions *)
-      List.map expanded ~f:(fun e -> { e with loc = expr.loc })
+      expand_matched_macro macro_env expand_macros_in_expr formal_params body
+        call_args expr.loc
     | None ->
       (* Not a macro - recursively expand in sub-expressions *)
       let expanded_subexprs =
-        List.map call_args ~f:(fun arg ->
-          match expand_macros_in_expr macro_env arg with
-          | [ single ] -> single
-          | multiple -> { content = List multiple; loc = arg.loc } )
+        expand_args_in_list macro_env expand_macros_in_expr call_args
       in
       [ { expr with
           content =
@@ -272,10 +312,7 @@ let rec expand_macros_in_expr
   (* Regular list - recursively expand in sub-expressions *)
   | List subexprs ->
     let expanded_subexprs =
-      List.map subexprs ~f:(fun sub ->
-        match expand_macros_in_expr macro_env sub with
-        | [ single ] -> single
-        | multiple -> { content = List multiple; loc = sub.loc } )
+      expand_args_in_list macro_env expand_macros_in_expr subexprs
     in
     [ { expr with content = List expanded_subexprs } ]
   (* Atoms - no expansion needed *)
@@ -296,7 +333,21 @@ let unfold_decl_def (macro_env : (string * (string list * expr loc list)) list)
             match arg.content with
             | Var x -> x
             | Symbol "..." -> "..." (* Allow ... as a symbol *)
-            | _ -> failwith "error: syntax declaration must contain variables" )
+            | Symbol s ->
+              raise
+                (MacroError
+                   ( InvalidMacroArgument
+                       (Printf.sprintf
+                          "macro argument '%s' must be a variable (start with \
+                           uppercase)"
+                          s )
+                   , arg.loc ) )
+            | List _ ->
+              raise
+                (MacroError
+                   ( InvalidMacroArgument
+                       "macro argument must be a variable, not a list"
+                   , arg.loc ) ) )
         in
         (macro_name, (var_args, body)) :: acc
       | _ -> acc )
@@ -362,7 +413,21 @@ let extract_macros (raw_exprs : Raw.t list) : macro_env =
           match arg.content with
           | Var x -> x
           | Symbol "..." -> "..." (* Allow ... as a symbol *)
-          | _ -> failwith "error: syntax declaration must contain variables" )
+          | Symbol s ->
+            raise
+              (MacroError
+                 ( InvalidMacroArgument
+                     (Printf.sprintf
+                        "macro argument '%s' must be a variable (start with \
+                         uppercase)"
+                        s )
+                 , arg.loc ) )
+          | List _ ->
+            raise
+              (MacroError
+                 ( InvalidMacroArgument
+                     "macro argument must be a variable, not a list"
+                 , arg.loc ) ) )
       in
       ((macro_name, (var_args, body)) :: env, acc)
     | _ -> (env, acc)

src/expr_err.ml

@@ -4,3 +4,10 @@ type expr_err =
   | InvalidBan of string
   | InvalidRaylist of string
   | InvalidDeclaration of string
+  | InvalidMacroArgument of string
+  | InvalidBanStructure of string
+  | CircularImport of string
+  | FileLoadError of
+      { filename : string
+      ; message : string
+      }

src/lexer.ml

@@ -21,7 +21,9 @@ let opposite_delimiter = function
   | ')' -> '('
   | ']' -> '['
   | '}' -> '{'
-  | c -> failwith (Printf.sprintf "Compiler error: '%c' is not a delimiter." c)
+  | _ ->
+    (* This case should be unreachable - if hit, it indicates a bug in the lexer *)
+    assert false
 
 let pop_delimiter sym (pos : Lexing.position) =
   match !delimiters_stack with