From cab9f10f62e2b32081932b9ece50581d3c38b8ab Mon Sep 17 00:00:00 2001
From: Kabir Samsi <samsikabir@gmail.com>
Date: Thu, 31 Oct 2024 21:59:28 -0400
Subject: [PATCH] WFQ Fixes

---
 progs/work_conserving/wfq_n_classes.sched |  2 +-
 rio/frontend/ast.ml                       |  2 +-
 rio/frontend/lexer.mll                    |  2 ++
 rio/frontend/parser.mly                   | 24 +++++++++++++----------
 rio/frontend/policy.ml                    |  4 ++--
 5 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/progs/work_conserving/wfq_n_classes.sched b/progs/work_conserving/wfq_n_classes.sched
index a01bfcc..fa0c1c6 100644
--- a/progs/work_conserving/wfq_n_classes.sched
+++ b/progs/work_conserving/wfq_n_classes.sched
@@ -1,5 +1,5 @@
 classes A, B, C;
 
-policy = wfq[fifo[A], fifo[B], fifo[C]][1, 2, 3];
+policy = wfq[(fifo[A], 1), (fifo[B], 2), (fifo[C], 3)];
 
 return policy
diff --git a/rio/frontend/ast.ml b/rio/frontend/ast.ml
index a96d653..5e0aaf1 100644
--- a/rio/frontend/ast.ml
+++ b/rio/frontend/ast.ml
@@ -15,7 +15,7 @@ type stream =
   (* Stream-To-Stream *)
   | RoundRobin of stream list
   | Strict of stream list
-  | WeightedFair of stream list * int list
+  | WeightedFair of (stream * int) list
   (* Non-Work Conserving *)
   | RateControlled of stream list
   | LeakyBucket of stream list * int * int
diff --git a/rio/frontend/lexer.mll b/rio/frontend/lexer.mll
index 24e38cf..f844cbf 100644
--- a/rio/frontend/lexer.mll
+++ b/rio/frontend/lexer.mll
@@ -16,6 +16,8 @@ rule token = parse
 | "="                   { EQUALS }
 | "["                   { LBRACKET }
 | "]"                   { RBRACKET }
+| "("                   { LPAREN }
+| ")"                   { RPAREN }
 | "return"              { RETURN }
 | "classes"             { CLASSES }
 | "union"               { UNION }
diff --git a/rio/frontend/parser.mly b/rio/frontend/parser.mly
index 196563f..a99f862 100644
--- a/rio/frontend/parser.mly
+++ b/rio/frontend/parser.mly
@@ -38,6 +38,8 @@
 %token EQUALS
 %token LBRACKET
 %token RBRACKET
+%token LPAREN
+%token RPAREN
 %token RETURN
 %token CLASSES
 %token UNION
@@ -71,14 +73,14 @@ set:
 
 /* Policies */
 policy:
-    | FIFO LBRACKET; pl = set; RBRACKET                                              { Fifo pl }
-    | EDF LBRACKET; pl = set; RBRACKET                                               { EarliestDeadline pl }
-    | SJN LBRACKET; pl = set; RBRACKET                                               { ShortestJobNext pl }
-    | SRTF LBRACKET; pl = set; RBRACKET                                              { ShortestRemaining pl }
-    | RR LBRACKET; pl = arglist; RBRACKET                                            { RoundRobin pl }
-    | STRICT LBRACKET; pl = arglist; RBRACKET                                        { Strict pl }
-    | WFQ LBRACKET; pl = arglist; RBRACKET; LBRACKET; wt = weight_list; RBRACKET     { WeightedFair (pl, wt) }
-    | RCSP LBRACKET; pl = arglist; RBRACKET                                          { RateControlled pl }
+    | FIFO LBRACKET; pl = set; RBRACKET                 { Fifo pl }
+    | EDF LBRACKET; pl = set; RBRACKET                  { EarliestDeadline pl }
+    | SJN LBRACKET; pl = set; RBRACKET                  { ShortestJobNext pl }
+    | SRTF LBRACKET; pl = set; RBRACKET                 { ShortestRemaining pl }
+    | RR LBRACKET; pl = arglist; RBRACKET               { RoundRobin pl }
+    | STRICT LBRACKET; pl = arglist; RBRACKET           { Strict pl }
+    | WFQ LBRACKET; pl = weighted_arglist; RBRACKET     { WeightedFair pl }
+    | RCSP LBRACKET; pl = arglist; RBRACKET             { RateControlled pl }
 
     | LEAKY LBRACKET; LBRACKET;
         pl = arglist; RBRACKET; COMMA;
@@ -99,8 +101,10 @@ setlist:
     | pl = separated_list(COMMA, set)               { pl }
 arglist:
     | pl = separated_list(COMMA, policy)            { pl }
-weight_list:
-    | pl = separated_list(COMMA, INT)               { pl }
+weighted_arglist:
+    | pl = separated_list(COMMA, weighted_arg)      { pl }
+weighted_arg:
+    | LPAREN; arg = separated_pair(policy, COMMA, INT); RPAREN      { arg }
 
 /* Declarations, assignments and returns */
 internalcomp :
diff --git a/rio/frontend/policy.ml b/rio/frontend/policy.ml
index c932dc9..fe0cd66 100644
--- a/rio/frontend/policy.ml
+++ b/rio/frontend/policy.ml
@@ -51,11 +51,11 @@ let rec sub cl st (p : Ast.policy) used =
   | Fifo p -> Fifo (sub_set cl st p used)
   | RoundRobin plst -> extract_subpol (RoundRobin (sub_plst cl st plst))
   | Strict plst -> extract_subpol (Strict (sub_plst cl st plst))
-  | WeightedFair (plst, wts) ->
+  | WeightedFair wplst ->
       extract_subpol
         (WeightedFair
            (sub_weighted_plst cl st
-              (List.combine plst (List.map float_of_int wts))))
+              (List.map (fun (x, y) -> (x, float_of_int y)) wplst)))
   | _ -> failwith "ERROR: unsupported policy"
 
 (* Look up any variables and substitute them in. *)