diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
index c3171f486..40c194efe 100644
--- a/lib/irb/ruby-lex.rb
+++ b/lib/irb/ruby-lex.rb
@@ -48,7 +48,7 @@ def self.compile_with_errors_suppressed(code, line_no: 1)
   end
 
   # io functions
-  def set_input(io, p = nil, context:, &block)
+  def set_input(io, context:, &block)
     @io = io
     if @io.respond_to?(:check_termination)
       @io.check_termination do |code|
@@ -116,22 +116,15 @@ def set_input(io, p = nil, context:, &block)
       end
     end
 
-    if p.respond_to?(:call)
-      @input = p
-    elsif block_given?
+    if block_given?
       @input = block
     else
       @input = Proc.new{@io.gets}
     end
   end
 
-  def set_prompt(p = nil, &block)
-    p = block if block_given?
-    if p.respond_to?(:call)
-      @prompt = p
-    else
-      @prompt = Proc.new{print p}
-    end
+  def set_prompt(&block)
+    @prompt = block
   end
 
   ERROR_TOKENS = [
diff --git a/test/irb/test_ruby_lex.rb b/test/irb/test_ruby_lex.rb
index 1c1ce92e6..8c8aa7407 100644
--- a/test/irb/test_ruby_lex.rb
+++ b/test/irb/test_ruby_lex.rb
@@ -62,7 +62,9 @@ def ruby_lex_for_lines(lines, local_variables: [])
 
       context = build_context(local_variables)
       io = proc{ lines.join("\n") }
-      ruby_lex.set_input(io, io, context: context)
+      ruby_lex.set_input(io, context: context) do
+        lines.join("\n")
+      end
       ruby_lex.lex(context)
       ruby_lex
     end