1 files changed, 7 insertions, 0 deletions
diff --git a/polymatrix/expression/mixins/fromstatemonad.py b/polymatrix/expression/mixins/fromstatemonad.py
index 216b160..a5970b1 100644
--- a/polymatrix/expression/mixins/fromstatemonad.py
+++ b/polymatrix/expression/mixins/fromstatemonad.py
@@ -25,12 +25,19 @@ class FromStateMonadMixin(ExpressionBaseMixin):
     @override
     def apply(self, state: ExpressionStateMixin) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
         state, expr = self.monad.apply(state)
+        # Case when monad wraps function 
+        #  f: ExpressionState -> (State, Expression)
         if isinstance(expr, Expression):
             return expr.underlying.apply(state)
 
+        # Case when monad wraps function 
+        #  f: ExpressionState -> (State, M)
+        # where M is a mixin class
         elif isinstance(expr, ExpressionBaseMixin):
             return expr.apply(state)
 
+        # Case when monad wraps function 
+        #  f: ExpressionState -> (State, PolyMatrix)
         elif isinstance(expr, PolyMatrixMixin):
             return state, expr