diff --git a/lib/pathway.rb b/lib/pathway.rb
index b0245f7..bf1b025 100644
--- a/lib/pathway.rb
+++ b/lib/pathway.rb
@@ -166,9 +166,9 @@ def set(callable, *args, to: @operation.result_key, **kwargs, &bl)
         end
 
         # Execute step and replace the current state completely
-        def map(callable)
+        def map(callable,...)
           bl = _callable(callable)
-          @result = @result.then(bl)
+          @result = @result.then { |state| bl.call(state,...) }
         end
 
         def around(execution_strategy, &dsl_block)
diff --git a/spec/plugins/base_spec.rb b/spec/plugins/base_spec.rb
index fff6375..ca08feb 100644
--- a/spec/plugins/base_spec.rb
+++ b/spec/plugins/base_spec.rb
@@ -12,6 +12,7 @@ class OperationWithSteps < Operation
 
         process do
           step :custom_validate
+          map  :add_misc
           set  :get_value
           set  :get_aux_value, to: :aux_value
           around(-> run, st { run.call if cond.call(st) }) do
@@ -34,6 +35,10 @@ def custom_validate(state)
           state[:params] = @validator.call(state)
         end
 
+        def add_misc(state)
+          State.new(self, state.to_h.merge(misc: -1))
+        end
+
         def get_value(state)
           @back_end.call(state[:params])
         end
@@ -112,6 +117,26 @@ def notify(state)
         end
       end
 
+      describe "#map" do
+        it "defines a step that replaces the current state" do
+          old_state = nil
+
+          expect(validator).to receive(:call) do |state|
+            expect(state.to_h).to_not include(:misc)
+
+            old_state = state
+            state[:input]
+          end
+
+          expect(notifier).to receive(:call) do |state|
+            expect(state.to_h).to include(misc: -1)
+            expect(state).to_not be_equal(old_state)
+          end
+
+          operation.call(input)
+        end
+      end
+
       describe "#set" do
         it "defines an updating step which sets the result key if no key is specified" do
           expect(back_end).to receive(:call).and_return(:SOME_VALUE)