Work on the dynamatic back end by adding type inference

Graphiti/Core/Basic.lean

@@ -75,6 +75,14 @@ meta def fromExpr? (e : Expr) : SimpM (Option Nat) :=
 meta def fromExpr?' (e : Expr) : SimpM (Option (Array Char)) :=
   getListLitOf? e getCharValue?
 
+def ofOption {ε α σ} (e : ε) : Option α → EStateM ε σ α
+| some o => pure o
+| none => throw e
+
+def ofOption' {ε α} (e : ε) : Option α → Except ε α
+| some o => pure o
+| none => throw e
+
 /--
 Reduce `toString 5` to `"5"`
 -/

Graphiti/Core/DotParser.lean

@@ -8,6 +8,7 @@ module
 
 public import Lean
 public import Graphiti.Core.ExprHigh
+public import Graphiti.Core.DynamaticTypes
 
 public section
 
@@ -269,49 +270,19 @@ def dotToExprHigh (d : Parser.DotGraph) : Except String (ExprHigh String String
         current_extra_args ← addOpt current_extra_args "taggers_num"
         current_extra_args ← addOpt current_extra_args "tagger_id"
 
-      -- Different bitwidths matter so we distinguish them with types: Unit vs. Bool vs. T
-      if typVal = "Branch" then
-        if keyStartsWith l "in" "in1:0" then
-          typVal := "branch Unit"
-        else if keyStartsWith l "in" "in1:1" then
-          typVal := "branch Bool"
-        -- Else, if the bitwidth is 32, then:
-        else typVal := "branch T"
-
-      -- Different bitwidths matter so we distinguish them with types: Unit vs. Bool vs. T
-      if typVal = "Fork" then
-        let [sizesIn] ← parseIOSizes l "in" |>.mapM translateSize
-          | throw "more that one input in fork"
-        typVal := s!"fork{keyArgNumbers l "out"} {sizesIn}"
-
-      if typVal = "Mux" then
+      if typVal == "Mux" || typVal == "Merge" then
         current_extra_args ← addOpt current_extra_args "delay"
-        if splitAndSearch l "in" "in2:0" then
-          typVal := s!"mux Unit"
-        else typVal := s!"mux T"
-
-      if typVal = "Merge" then
-        current_extra_args ← addOpt current_extra_args "delay"
-        if splitAndSearch l "in" "in0:0" then
-          typVal := s!"merge{keyArgNumbers l "in"} Unit"
-        else typVal := s!"merge{keyArgNumbers l "in"} T"
 
       if typVal = "Entry" then
         current_extra_args ← addOpt current_extra_args "control"
 
-
       if typVal = "Constant" then
         let constVal ← keyArg l "value" |> Parser.hexParser.run
-        typVal := s!"constant {constVal}"
-        current_extra_args ← addOpt current_extra_args "value"
-
-      if typVal = "Sink" then
-        if splitAndSearch l "in" "in1:0" then
-          typVal := s!"sink Unit 1"
-        else if splitAndSearch l "in" "in1:1" then
-          typVal := s!"sink Bool 1"
+        if keyStartsWith l "out" "out1:1" then
+          typVal := s!"constantBool"
         else
-          typVal := s!"sink T 1"
+          typVal := s!"constantNat"
+        current_extra_args ← addOpt current_extra_args "value"
 
       if typVal = "Operator" then
         if splitAndSearch l "op" "mc_store_op" || splitAndSearch l "op" "mc_load_op" then
@@ -325,25 +296,29 @@ def dotToExprHigh (d : Parser.DotGraph) : Except String (ExprHigh String String
         current_extra_args ← add current_extra_args "op"
 
         if splitAndSearch l "op" "mc_load_op" then
-          typVal := s!"load T T"
-        else if splitAndSearch l "op" "cast" then
-          typVal := s!"pure T Bool"
+          typVal := s!"load"
         else
           let sizesIn ← parseIOSizes l "in" |>.mapM translateSize
           let sizesOut ← parseIOSizes l "out" |>.mapM translateSize
-          typVal := s!"operator{keyArgNumbers l "in"} {" ".intercalate sizesIn} {" ".intercalate sizesOut} {keyArg l "op"}"
-
-       -- portId= 0, offset= 0  -- if mc_store_op and mc_load_op
+          typVal := s!"operator{keyArgNumbers l "in"}"
 
       if typVal = "MC" then
         let sizesIn ← parseIOSizes l "in" |>.mapM translateSize
         let sizesOut ← parseIOSizes l "out" |>.filter (·.trim.endsWith "*e" |> not) |>.mapM translateSize
-        typVal := s!"operator{keyArgNumbers l "in"} {" ".intercalate sizesIn} {" ".intercalate sizesOut} MC"
+        typVal := s!"operator{keyArgNumbers l "in"}"
         current_extra_args ← addOpt current_extra_args "memory"
         current_extra_args ← addOpt current_extra_args "bbcount"
         current_extra_args ← addOpt current_extra_args "ldcount"
         current_extra_args ← addOpt current_extra_args "stcount"
 
+      if typVal == "Fork" then
+        typVal := s!"fork{keyArgNumbers l "out"}"
+      if typVal == "Merge" then
+        typVal := s!"merge{keyArgNumbers l "in"}"
+
+      unless "merge".isPrefixOf typVal || "fork".isPrefixOf typVal || "constant".isPrefixOf typVal || "operator".isPrefixOf typVal do
+        typVal := dynamaticToGraphiti typVal
+
       let cluster := l.find? (·.key = "cluster") |>.getD ⟨"cluster", "false"⟩
       let .ok clusterB := Parser.parseBool.run cluster.value.trim
         | throw s!"{s}: cluster could not be parsed"

Graphiti/Core/DynamaticPrinter.lean

@@ -8,6 +8,8 @@ module
 
 public import Graphiti.Core.Rewriter
 public import Graphiti.Core.TypeExpr
+public import Graphiti.Core.DynamaticTypes
+public import Graphiti.Core.WellTyped
 
 public section
 
@@ -107,54 +109,6 @@ def translateTypes  (key : String) : Option String × String × String × List (
 def removeLetter (ch : Char) (s : String) : String :=
   String.mk (s.toList.filter (λ c => c ≠ ch))
 
-def returnNatInstring (s : String) : Option Nat :=
-  -- Convert the string to a list of characters
-  let chars := s.toList
-  let result := List.foldl (λ acc c =>
-      if c.isDigit then
-        acc * 10 + (Char.toNat c - Char.toNat '0')
-      else
-        acc) 0 chars
-  -- If no non-digit character was encountered, return the result
-  -- if result = 0 then
-  --   if s.isEmpty then some 0 else none
-  -- else
-  some result
-
-def incrementDefinitionPortIdx (s direction: String) : String :=
-  -- Split the string by spaces into individual parts (like "out0:32")
-  let parts := s.splitOn " "
-  -- Map over each part, incrementing the number after "out"
-  let updatedParts := parts.map (λ part =>
-    match part.splitOn ":" with
-    | [pref, num] =>
-      match (returnNatInstring pref) with
-      | some n =>
-        -- Increment the number found
-        let incremented := n + 1
-        -- Reconstruct the string with the incremented number
-        direction ++ Nat.repr incremented ++ ":" ++ (num)
-      | none => part -- If no number is found, keep the part unchanged
-    | _ => part -- If the part doesn't have ":" or a valid number, keep it unchanged
-  )
-  -- Join the updated parts into a single string with spaces
-  String.intercalate " " updatedParts
-
--- #eval incrementDefinitionPortIdx "out1:32" "out"  --out1:324 out2:32 out3:32" "out"  -- Output: "out1:32 out4:32 out3:32"
-
--- #eval "out132".splitOn ":"
-
-def incrementConnectionPortIdx (s direction: String) : String :=
-   match returnNatInstring s with
-  | some n =>
-    let incremented := n + 1
-    -- Convert incremented number to a string and concatenate with the direction part
-    let incrementedStr := Nat.repr incremented
-    direction ++ incrementedStr
-  | none => s  -- If no number is found, return the original string
-
--- #eval incrementConnectionPortIdx "out33" "out"
-
 -- Function became messy...
 def formatOptions : List (String × String) → String
 | x :: l => l.foldl
@@ -167,56 +121,43 @@ def formatOptions : List (String × String) → String
      s!", {x.1} = {v2_}")
 | [] => ""
 
-def extractStandardType (s : String) : String :=
-  let parts := s.splitOn " "
-  parts[0]!
-
-def capitalizeFirstChar (s : String) : String :=
-  match s.get? 0 with
-  | none   => s  -- If the string is empty, return it as is
-  | some c =>
-    let newChar := if 'a' ≤ c ∧ c ≤ 'z' then
-                     Char.ofNat (c.toNat - ('a'.toNat - 'A'.toNat))
-                   else
-                     c
-    newChar.toString ++ s.drop 1
-
--- Join is taken in Dynamatic so rename to Concat
-def RenameJoinToConcat (s : String) : String :=
-  if String.isPrefixOf "join" s then
-    "Concat"
-  else
-    s  -- Otherwise, return the original string
-
-def fixComponentNames (s : String) : String :=
-  String.intercalate "_" (s.splitOn "__")
+def inferTypeInPortMap (t : TypeUF) (p : PortMap String (InternalPort String)) (sn : String × Nat) : Except String (PortMap String TypeExpr) :=
+  p.foldlM (λ st k v => do
+      let tc ← toTypeConstraint sn k.name
+      let concr ← ofOption' "could not find concr" <| t.findConcr tc
+      return st.cons k concr
+    ) ∅
+
+def inferTypeInPortMapping (t : TypeUF) (p : PortMapping String) (sn : String × Nat) : Except String (PortMap String TypeExpr × PortMap String TypeExpr) := do
+  let inp ← inferTypeInPortMap t p.input sn
+  let out ← inferTypeInPortMap t p.output sn
+  return (inp, out)
+
+def toPortList (typs : PortMap String TypeExpr) : String :=
+  typs.foldl (λ s k v => s ++ s!"{removeLetter 'p' k.name}:{TypeExpr.Parser.getSize v} ") ""
 
 --fmt.1: Type
 --fmt.2.1 and fmt.2.2.1: Input and output attributes
 --fmt.2.2.2: Additional options.
-def dynamaticString (a: ExprHigh String (String × Nat)) (m : AssocList String (AssocList String String)): Option String := do
-  -- let instances :=
-  --   a.modules.foldl (λ s inst mod => s ++ s!"\n {inst} [mod = \"{mod}\"];") ""
-  let a ← a.normaliseNames
+def dynamaticString (a: ExprHigh String (String × Nat)) (t : TypeUF) (m : AssocList String (AssocList String String)): Except String String := do
+  let a ← ofOption' "could not normalise names" a.normaliseNames
   let modules ←
     a.modules.foldlM
       (λ s k v => do
         -- search for the type of the passed node in interfaceTypes
-        let fmt := translateTypes v.snd.1
+        let typeName ← ofOption' s!"could not convert graphiti to dynamatic name for {k}/{v.2.1}" <| graphitiToDynamatic v.2.1
         match m.find? k with
         | some input_fmt =>
           -- If the node is found to be coming from the input,
           -- retrieve its attributes from what we saved and bypass it
           -- without looking for it in interfaceTypes
-        return (RenameJoinToConcat s) ++ s!"\"{k}\" [type = \"{capitalizeFirstChar (extractStandardType (fmt.1.getD v.snd.1))}\"{formatOptions input_fmt.toList}];\n"
-        --return s ++ s!"\"{k}\" [type = \"{fmt.1.getD v.snd}\"{formatOptions input_fmt.toList}];\n"
+          return s ++ s!"\"{k}\" [type = \"{typeName}\"{formatOptions input_fmt.toList}];\n"
         | none =>
+          let typs ← inferTypeInPortMapping t v.1 v.2
           -- If this is a new node, then we sue `fmt` to correctly add the right
           -- arguments from what is given in interfaceTypes.  We should never be generating constructs like MC, so
           -- this shouldn't be a problem.
-        return (RenameJoinToConcat s) ++ s!"\"{k}\" [type = \"{capitalizeFirstChar (extractStandardType (fmt.1.getD v.snd.1))}\", in = \"{removeLetter 'p' fmt.2.1}\", out = \"{fmt.2.2.1}\"{formatOptions fmt.2.2.2}];\n"
-        --return s ++ s!"\"{k}\" [type = \"{fmt.1.getD v.snd}\", in = \"{removeLetter 'p' fmt.2.1}\", out = \" {fmt.2.2.1} \"{formatOptions fmt.2.2.2}];\n"
-
+          return s ++ s!"\"{k}\" [type = \"{typeName}\", in = \"{toPortList typs.1}\", out = \"{toPortList typs.2}\"];\n"
       ) ""
   let connections :=
     a.connections.foldl
@@ -225,7 +166,7 @@ def dynamaticString (a: ExprHigh String (String × Nat)) (m : AssocList String (
                     ++ s!"[from = \"{oport.name}\","
                     ++ s!" to = \"{removeLetter 'p' iport.name}\" "
                     ++ "];") ""
-  s!"Digraph G \{
+  .ok s!"Digraph G \{
 {modules}
 {connections}
 }"

Graphiti/Core/DynamaticTypes.lean

@@ -0,0 +1,62 @@
+/-
+Copyright (c) 2025 VCA Lab, EPFL. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Yann Herklotz
+-/
+
+module
+
+public import Graphiti.Core.AssocList
+
+public section
+
+open Batteries (AssocList)
+
+namespace Graphiti
+
+def graphitiPrefix := "_graphiti_"
+
+/--
+Mapping from Graphiti names to Dynamatic names.
+
+TODO: This needs to be a bit less ad-hoc, i.e. a true 2-way mapping with port information from Dynamatic.
+-/
+def dynamatic_types : AssocList String String :=
+  [ ("join", "Concat")
+  , ("split", "Split")
+  , ("branch", "Branch")
+  , ("fork2", "Fork")
+  , ("fork3", "Fork")
+  , ("fork4", "Fork")
+  , ("fork5", "Fork")
+  , ("fork", "Fork")
+  , ("merge", "Merge")
+  , ("merge2", "Merge")
+  , ("operator1", "Operator")
+  , ("operator2", "Operator")
+  , ("operator3", "Operator")
+  , ("operator4", "Operator")
+  , ("operator5", "Operator")
+  , ("mux", "Mux")
+  , ("input", "Entry")
+  , ("output", "Exit")
+  , ("sink", "Sink")
+  , ("constantNat", "Constant")
+  , ("constantBool", "Constant")
+  , ("initBool", "Init")
+  , ("tagger_untagger_val", "TaggerUntagger")
+  , ("load", "Operator")
+  ].toAssocList
+
+def dynamaticToGraphiti (s : String) : String :=
+  match dynamatic_types.inverse.find? s with
+  | some s' => s'
+  | none => graphitiPrefix ++ s
+
+def graphitiToDynamatic (s : String) : Option String :=
+  match dynamatic_types.find? s with
+  | some s' => some s'
+  | none =>
+    if graphitiPrefix.isPrefixOf s then some (s.drop graphitiPrefix.length) else none
+
+end Graphiti

Graphiti/Core/Rewriter.lean

@@ -121,14 +121,6 @@ variable {Ident Typ}
 variable [Inhabited Ident]
 variable [Inhabited Typ]
 
-def ofOption {ε α σ} (e : ε) : Option α → EStateM ε σ α
-| some o => pure o
-| none => throw e
-
-def ofOption' {ε α} (e : ε) : Option α → Except ε α
-| some o => pure o
-| none => throw e
-
 def liftError {α σ} : Except String α → EStateM RewriteError σ α
 | .ok o => pure o
 | .error s => throw (.error s)

Graphiti/Core/TypeExpr.lean

@@ -22,6 +22,7 @@ inductive TypeExpr where
 | bool
 | tag
 | unit
+| var (n : Nat)
 | pair (left right : TypeExpr)
 deriving Repr, DecidableEq, Inhabited, Hashable
 
@@ -32,6 +33,7 @@ def toBlueSpec : TypeExpr → String
 | .tag => "Token"
 | .bool => "Bool"
 | .unit => "Void"
+| .var n => s!"Var#({n})"
 | .pair left right =>
   s!"Tuple2#({toBlueSpec left}, {toBlueSpec right})"
 
@@ -77,6 +79,7 @@ noncomputable def TypeExpr.denote : TypeExpr → Type
 | tag => Unit
 | bool => Bool
 | unit => Unit
+| var n => Unit
 | pair t1 t2 => t1.denote × t2.denote
 
 def ValExpr.type : ValExpr → TypeExpr
@@ -158,6 +161,7 @@ def toString' : TypeExpr → String
   | TypeExpr.tag => "TagT" -- Unclear how we want to display TagT at the end?
   | TypeExpr.bool => "Bool"
   | TypeExpr.unit => "Unit"
+  | TypeExpr.var n => s!"Var({n})"
   | TypeExpr.pair left right =>
     let leftStr := toString' left
     let rightStr :=  toString' right
@@ -168,6 +172,7 @@ def getSize: TypeExpr → Int
   | TypeExpr.tag => 0
   | TypeExpr.bool => 1
   | TypeExpr.unit => 0
+  | TypeExpr.var n => 0
   | TypeExpr.pair left right =>
     let l := getSize left
     let r :=  getSize right
@@ -201,6 +206,7 @@ def flatten_type (t : TypeExpr) : List TypeExpr :=
   | TypeExpr.tag => [t]
   | TypeExpr.bool => [t]
   | TypeExpr.unit => [t]
+  | TypeExpr.var _ => [t]
   | TypeExpr.pair left right =>
     flatten_type left ++ flatten_type right