diff --git a/Sources/FirebladeECS/FSM.swift b/Sources/FirebladeECS/FSM.swift
index 01b22b4..bc75c19 100644
--- a/Sources/FirebladeECS/FSM.swift
+++ b/Sources/FirebladeECS/FSM.swift
@@ -122,14 +122,14 @@ extension ComponentSingletonProvider: ComponentProvider {
 
 /// This component provider calls a function to get the component instance. The function must
 /// return a single component of the appropriate type.
-public class DynamicComponentProvider {
+public class DynamicComponentProvider<C: Component> {
     /// Wrapper for closure to make it hashable via ObjectIdentifier
     public class Closure {
-        let provideComponent: () -> Component
+        let provideComponent: () -> C
 
         /// Initializer
         /// - Parameter provideComponent: Swift closure returning component of the appropriate type
-        public init(provideComponent: @escaping () -> Component) {
+        public init(provideComponent: @escaping () -> C) {
             self.provideComponent = provideComponent
         }
     }
@@ -231,12 +231,12 @@ extension EntityState {
 
     /// Creates a mapping for the component type to a method call.
     /// A DynamicComponentProvider is used for the mapping.
-    /// - Parameter method: The method to return the component instance
+    /// - Parameter closure: The Closure instance to return the component instance
     /// - Returns: This EntityState, so more modifications can be applied
     @inline(__always)
     @discardableResult
-    public func addMethod<C: ComponentInitializable>(type: C.Type, closure: DynamicComponentProvider.Closure) -> Self {
-        add(type).withMethod(closure)
+    public func addMethod<C: ComponentInitializable>(closure: DynamicComponentProvider<C>.Closure) -> Self {
+        add(C.self).withMethod(closure)
         return self
     }
 
@@ -308,10 +308,10 @@ public class StateComponentMapping {
 
     /// Creates a mapping for the component type to a method call. A
     /// DynamicComponentProvider is used for the mapping.
-    /// - Parameter method: The method to return the component instance
+    /// - Parameter closure: The Closure instance to return the component instance
     /// - Returns: This ComponentMapping, so more modifications can be applied
     @discardableResult
-    public func withMethod(_ closure: DynamicComponentProvider.Closure) -> Self {
+    public func withMethod<C: Component>(_ closure: DynamicComponentProvider<C>.Closure) -> Self {
         setProvider(DynamicComponentProvider(closure: closure))
         return self
     }
diff --git a/Tests/FirebladeECSTests/FSMTests.swift b/Tests/FirebladeECSTests/FSMTests.swift
index f31c372..28fa552 100644
--- a/Tests/FirebladeECSTests/FSMTests.swift
+++ b/Tests/FirebladeECSTests/FSMTests.swift
@@ -213,7 +213,7 @@ class EntityStateTests: XCTestCase {
         state.add(MockComponent.self).withMethod(dynamickProvider)
         let provider = state.providers[MockComponent.identifier]
         XCTAssertNotNil(provider)
-        XCTAssertTrue(provider is DynamicComponentProvider?)
+        XCTAssertTrue(provider is DynamicComponentProvider<MockComponent>?)
         XCTAssertTrue(provider?.getComponent() is MockComponent)
     }
 
@@ -243,7 +243,7 @@ class EntityStateTests: XCTestCase {
         state.add(MockComponent.self).withMethod(.init { MockComponent() })
         let provider = state.get(MockComponent.self)
         XCTAssertNotNil(provider)
-        XCTAssertTrue(provider is DynamicComponentProvider?)
+        XCTAssertTrue(provider is DynamicComponentProvider<MockComponent>?)
     }
 
     func testGetReturnsTypeProvider() {
@@ -293,8 +293,8 @@ class EntityStateTests: XCTestCase {
 
     func testAddMethodCreatesMappingAndSetsDynamicProviderForType() {
         let component = MockComponent()
-        state.addMethod(type: MockComponent.self, closure: .init { component })
-        XCTAssertTrue(state.get(MockComponent.self) is DynamicComponentProvider?)
+        state.addMethod(closure: .init { component })
+        XCTAssertTrue(state.get(MockComponent.self) is DynamicComponentProvider<MockComponent>?)
         XCTAssertTrue(state.get(MockComponent.self)?.getComponent() === component)
     }