fix: ensure correct callback behavior with reusable states

Chickensoft.LogicBlocks.Example/VendingMachine.cs

@@ -17,7 +17,7 @@ public abstract record State(Context Context) : StateLogic(Context) {
     public record Idle : State,
       IGet<Input.SelectionEntered>, IGet<Input.PaymentReceived> {
       public Idle(Context context) : base(context) {
-        context.OnEnter<Idle>((previous) => context.Output(
+        OnEnter<Idle>((previous) => context.Output(
           new Output.ClearTransactionTimeOutTimer()
         ));
       }
@@ -55,7 +55,7 @@ public abstract record TransactionActive : State,
         Price = price;
         AmountReceived = amountReceived;
 
-        Context.OnEnter<TransactionActive>(
+        OnEnter<TransactionActive>(
          (previous) => Context.Output(
            new Output.RestartTransactionTimeOutTimer()
          )
@@ -97,7 +97,7 @@ public record Started : TransactionActive,
         public Started(
           Context context, ItemType type, int price, int amountReceived
         ) : base(context, type, price, amountReceived) {
-          context.OnEnter<Started>(
+          OnEnter<Started>(
             (previous) => context.Output(new Output.TransactionStarted())
           );
         }
@@ -128,7 +128,7 @@ public record Vending : State, IGet<Input.VendingCompleted> {
         Type = type;
         Price = price;
 
-        context.OnEnter<Vending>(
+        OnEnter<Vending>(
           (previous) => Context.Output(new Output.BeginVending())
         );
       }

Chickensoft.LogicBlocks.Generator.Tests/GeneratorTest.cs

@@ -11,7 +11,7 @@ public class GeneratorTest {
 
     result.Outputs["ToasterOven.puml.g.cs"].ShouldBe("""
     @startuml ToasterOven
-    state "ToasterOven" as State {
+    state "ToasterOven State" as Chickensoft_LogicBlocks_Generator_Tests_ToasterOven_State {
       state "Heating" as Chickensoft_LogicBlocks_Generator_Tests_ToasterOven_State_Heating {
         state "Toasting" as Chickensoft_LogicBlocks_Generator_Tests_ToasterOven_State_Toasting {
           Chickensoft_LogicBlocks_Generator_Tests_ToasterOven_State_Toasting : OnEnter → SetTimer

Chickensoft.LogicBlocks.Generator.Tests/test_cases/Heater.cs

@@ -52,26 +52,25 @@ public record Heating : State,
       ) {
         var tempSensor = context.Get<ITemperatureSensor>();
 
-        context.OnEnter<Heating>(
+        OnEnter<Heating>(
           (previous) => tempSensor.OnTemperatureChanged += OnTemperatureChanged
         );
 
-        context.OnExit<Heating>(
+        OnExit<Heating>(
           (next) => tempSensor.OnTemperatureChanged -= OnTemperatureChanged
         );
       }
 
-      State IGet<Input.TurnOff>.On(Input.TurnOff input)
-        => new Off(Context, TargetTemp);
+      public State On(Input.TurnOff input) => new Off(Context, TargetTemp);
 
-      State IGet<Input.AirTempSensorChanged>.On(
-        Input.AirTempSensorChanged input
-      ) => input.AirTemp >= TargetTemp
-        ? new Idle(Context, TargetTemp)
-        : this;
+      public State On(Input.AirTempSensorChanged input) =>
+        input.AirTemp >= TargetTemp
+          ? new Idle(Context, TargetTemp)
+          : this;
 
-      State IGet<Input.TargetTempChanged>.On(Input.TargetTempChanged input)
-        => this with { TargetTemp = input.Temp };
+      public State On(Input.TargetTempChanged input) => this with {
+        TargetTemp = input.Temp
+      };
 
       private void OnTemperatureChanged(double airTemp) {
         Context.Input(new Input.AirTempSensorChanged(airTemp));

Chickensoft.LogicBlocks.Generator.Tests/test_cases/ToasterOven.cs

@@ -15,10 +15,10 @@ public record StartToasting(int ToastColor) : Input;
   public abstract record State(Context Context) : StateLogic(Context) {
     public record Heating : State, IGet<Input.OpenDoor> {
       public Heating(Context context) : base(context) {
-        Context.OnEnter<Heating>(
+        OnEnter<Heating>(
           (previous) => Context.Output(new Output.TurnHeaterOn())
         );
-        Context.OnExit<Heating>(
+        OnExit<Heating>(
           (next) => Context.Output(new Output.TurnHeaterOff())
         );
       }
@@ -32,10 +32,10 @@ public record Toasting : Heating, IGet<Input.StartBaking> {
       public Toasting(Context context, int toastColor) : base(context) {
         ToastColor = toastColor;
 
-        Context.OnEnter<Toasting>(
+        OnEnter<Toasting>(
           (previous) => Context.Output(new Output.SetTimer(ToastColor))
         );
-        Context.OnExit<Toasting>(
+        OnExit<Toasting>(
           (next) => Context.Output(new Output.ResetTimer())
         );
       }
@@ -51,10 +51,10 @@ public record Baking : Heating, IGet<Input.StartToasting> {
       public Baking(Context context, int temperature) : base(context) {
         Temperature = temperature;
 
-        Context.OnEnter<Baking>(
+        OnEnter<Baking>(
           (previous) => Context.Output(new Output.SetTemperature(Temperature))
         );
-        Context.OnExit<Baking>(
+        OnExit<Baking>(
           (next) => Context.Output(new Output.SetTemperature(0))
         );
       }
@@ -66,10 +66,10 @@ public record Baking : Heating, IGet<Input.StartToasting> {
 
     public record DoorOpen : State, IGet<Input.CloseDoor> {
       public DoorOpen(Context context) : base(context) {
-        Context.OnEnter<DoorOpen>(
+        OnEnter<DoorOpen>(
           (previous) => Context.Output(new Output.TurnLampOn())
         );
-        Context.OnExit<DoorOpen>(
+        OnExit<DoorOpen>(
           (next) => Context.Output(new Output.TurnLampOff())
         );
       }

Chickensoft.LogicBlocks.Generator.Tests/test_cases/partial_split_across_files/PartialLogic1.cs

@@ -16,5 +16,6 @@ public abstract record Output {
     public record OutputA : Output;
     public record OutputEnterA : Output;
     public record OutputExitA : Output;
+    public record OutputSomething : Output;
   }
 }

Chickensoft.LogicBlocks.Generator.Tests/test_cases/partial_split_across_files/PartialLogic2.cs

@@ -5,10 +5,10 @@ public partial class PartialLogic :
   public abstract partial record State : StateLogic {
     public partial record A : State, IGet<Input.One> {
       public A(Context context) : base(context) {
-        Context.OnEnter<A>(
+        OnEnter<A>(
           (previous) => Context.Output(new Output.OutputEnterA())
         );
-        Context.OnExit<A>(
+        OnExit<A>(
           (next) => Context.Output(new Output.OutputExitA())
         );
       }

Chickensoft.LogicBlocks.Generator.Tests/test_cases/partial_split_across_files/PartialLogic3.cs

@@ -8,6 +8,8 @@ public partial record A : State, IGet<Input.One> {
         Context.Output(new Output.OutputA());
         return new B(Context);
       }
+
+      public void DoSomething() => Context.Output(new Output.OutputSomething());
     }
   }
 }

Chickensoft.LogicBlocks.Generator/Chickensoft.LogicBlocks.Generator.csproj

@@ -11,7 +11,7 @@
     <NoWarn>NU5128</NoWarn>
 
     <Title>LogicBlocks Generator</Title>
-    <Version>1.3.0</Version>
+    <Version>1.4.0</Version>
     <Description></Description>
     <Copyright>© 2023 Chickensoft Games</Copyright>
     <Authors>Chickensoft</Authors>

Chickensoft.LogicBlocks.Generator/src/LogicBlocksGenerator.cs

@@ -3,6 +3,7 @@ namespace Chickensoft.LogicBlocks.Generator;
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Text;
@@ -30,7 +31,7 @@ public class LogicBlocksGenerator :
     // the source generator process is started by running `dotnet build` in
     // the project consuming the source generator
     //
-    // Debugger.Launch();
+    Debugger.Launch();
 
     // Add post initialization sources
     // (source code that is always generated regardless)
@@ -235,20 +236,22 @@ CancellationToken token
           member.Name == Constants.LOGIC_BLOCK_GET_INITIAL_STATE
       );
 
-    string? initialStateId = null;
+    HashSet<string> initialStateIds = new();
 
     if (
       getInitialStateMethod is IMethodSymbol initialStateMethod &&
-      initialStateMethod
-        .DeclaringSyntaxReferences.FirstOrDefault()?.GetSyntax(token) is
-      SyntaxNode initialStateMethodNode
+      initialStateMethod.DeclaringSyntaxReferences.Select(
+        (syntaxRef) => syntaxRef.GetSyntax(token)
+      ).OfType<MethodDeclarationSyntax>() is
+        IEnumerable<MethodDeclarationSyntax> initialStateMethodSyntaxes
     ) {
-      var initialStateVisitor = new ReturnTypeVisitor(
-        model, token, CodeService, stateBaseType
-      );
-      initialStateVisitor.Visit(initialStateMethodNode);
-      initialStateId = initialStateVisitor.ReturnTypes.FirstOrDefault();
-      Log.Print($"Initial state type: {initialStateId}");
+      foreach (var initialStateMethodSyntax in initialStateMethodSyntaxes) {
+        var initialStateVisitor = new ReturnTypeVisitor(
+          model, token, CodeService, stateBaseType
+        );
+        initialStateVisitor.Visit(initialStateMethodSyntax);
+        initialStateIds.UnionWith(initialStateVisitor.ReturnTypes);
+      }
     }
 
     // Convert the subtypes into a graph by recursively building the graph
@@ -302,7 +305,7 @@ SyntaxNode initialStateMethodNode
       FilePath: destFile,
       Id: CodeService.GetNameFullyQualified(symbol, symbol.Name),
       Name: symbol.Name,
-      InitialStateId: initialStateId,
+      InitialStateIds: initialStateIds.ToImmutableHashSet(),
       Graph: root,
       Inputs: inputs.ToImmutableDictionary(),
       Outputs: outputs.ToImmutableDictionary(),
@@ -342,18 +345,21 @@ CancellationToken token
 
     transitions.Sort();
 
-    var initialStateString = implementation.InitialStateId != null
-      ? "[*] --> " +
-        $"{implementation.StatesById[implementation.InitialStateId].UmlId}"
-      : "";
+    var initialStates = new List<string>();
+
+    foreach (var initialStateId in implementation.InitialStateIds) {
+      initialStates.Add(
+        "[*] --> " + implementation.StatesById[initialStateId].UmlId
+      );
+    }
 
     var text = Format($"""
     @startuml {implementation.Name}
     {WriteGraph(implementation.Graph, implementation, 0)}
 
     {transitions}
 
-    {initialStateString}
+    {initialStates}
     @enduml
     """);
 
@@ -379,15 +385,15 @@ int t
     if (isMultilineState) {
       if (isRoot) {
         lines.Add(
-          $"{Tab(t)}state \"{impl.Name}\" as {graph.Name} {{"
+          $"{Tab(t)}state \"{impl.Name} State\" as {graph.UmlId} {{"
         );
       }
       else {
         lines.Add($"{Tab(t)}state \"{graph.Name}\" as {graph.UmlId} {{");
       }
     }
     else if (isRoot) {
-      lines.Add($"{Tab(t)}state \"{impl.Name} {graph.Name}\" as {graph.Name}");
+      lines.Add($"{Tab(t)}state \"{impl.Name} State\" as {graph.UmlId}");
     }
     else {
       lines.Add($"{Tab(t)}state \"{graph.Name}\" as {graph.UmlId}");
@@ -443,11 +449,11 @@ int t
   }
 
   public StatesAndOutputs GetStatesAndOutputs(
-    INamedTypeSymbol type,
-    SemanticModel model,
-    CancellationToken token,
-    INamedTypeSymbol stateBaseType
-  ) {
+      INamedTypeSymbol type,
+      SemanticModel model,
+      CancellationToken token,
+      INamedTypeSymbol stateBaseType
+    ) {
     // type is the state type
 
     var inputToStatesBuilder = ImmutableDictionary
@@ -477,12 +483,6 @@ Constants.LOGIC_BLOCK_INPUT_INTERFACE_ID or
       .SelectMany(syntaxNode => syntaxNode.ChildNodes())
       .OfType<ConstructorDeclarationSyntax>().ToList();
 
-    var handledInputInterfaceSyntaxes = handledInputInterfaces
-      .SelectMany(
-        interfaceType => interfaceType.DeclaringSyntaxReferences
-          .Select(syntaxRef => syntaxRef.GetSyntax(token))
-      );
-
     var inputHandlerMethods = new List<MethodDeclarationSyntax>();
 
     var outputVisitor = new OutputVisitor(
@@ -507,38 +507,46 @@ Constants.LOGIC_BLOCK_INPUT_INTERFACE_ID or
       if (implementation is not IMethodSymbol methodSymbol) {
         continue;
       }
-      var handlerMethodSyntax = methodSymbol
+
+      var handlerMethodSyntaxes = methodSymbol
         .DeclaringSyntaxReferences
-        .FirstOrDefault()?
-        .GetSyntax(token) as MethodDeclarationSyntax;
-      if (handlerMethodSyntax is not MethodDeclarationSyntax methodSyntax) {
+        .Select(syntaxRef => syntaxRef.GetSyntax(token))
+        .OfType<MethodDeclarationSyntax>()
+        .ToImmutableArray();
+
+      if (handlerMethodSyntaxes.Length == 0) {
         continue;
       }
-      inputHandlerMethods.Add(methodSyntax);
-      var inputId = CodeService.GetNameFullyQualifiedWithoutGenerics(
-        inputType, inputType.Name
-      );
-      var outputContext = OutputContexts.OnInput(inputType.Name);
-      var returnTypeVisitor = new ReturnTypeVisitor(
-        model, token, CodeService, stateBaseType
-      );
-      outputVisitor = new OutputVisitor(
-        model, token, CodeService, outputContext
-      );
 
-      returnTypeVisitor.Visit(methodSyntax);
-      outputVisitor.Visit(methodSyntax);
+      foreach (var methodSyntax in handlerMethodSyntaxes) {
+        inputHandlerMethods.Add(methodSyntax);
+        var inputId = CodeService.GetNameFullyQualifiedWithoutGenerics(
+          inputType, inputType.Name
+        );
+        var outputContext = OutputContexts.OnInput(inputType.Name);
+        var modelForSyntax =
+          model.Compilation.GetSemanticModel(methodSyntax.SyntaxTree);
+        var returnTypeVisitor = new ReturnTypeVisitor(
+          modelForSyntax, token, CodeService, stateBaseType
+        );
+        outputVisitor = new OutputVisitor(
+          modelForSyntax, token, CodeService, outputContext
+        );
+
+        returnTypeVisitor.Visit(methodSyntax);
+        outputVisitor.Visit(methodSyntax);
 
-      if (outputVisitor.OutputTypes.ContainsKey(outputContext)) {
-        outputsBuilder.Add(
-          outputContext, outputVisitor.OutputTypes[outputContext]
+        if (outputVisitor.OutputTypes.ContainsKey(outputContext)) {
+          outputsBuilder.Add(
+            outputContext, outputVisitor.OutputTypes[outputContext]
+          );
+        }
+
+        inputToStatesBuilder.Add(
+          inputId,
+          returnTypeVisitor.ReturnTypes
         );
       }
-
-      inputToStatesBuilder.Add(
-        inputId,
-        returnTypeVisitor.ReturnTypes
-      );
     }
 
     // find methods on type that aren't input handlers or constructors
@@ -553,8 +561,12 @@ Constants.LOGIC_BLOCK_INPUT_INTERFACE_ID or
       Log.Print("Examining method: " + otherMethod.Identifier.Text);
       var outputContext = OutputContexts.Method(otherMethod.Identifier.Text);
 
+      var modelForSyntax = model.Compilation.GetSemanticModel(
+        otherMethod.SyntaxTree
+      );
+
       outputVisitor = new OutputVisitor(
-        model, token, CodeService, outputContext
+        modelForSyntax, token, CodeService, outputContext
       );
       outputVisitor.Visit(otherMethod);