Skip to content

Commit 51e939b

Browse files
mameclaude
andcommitted
Add Struct.new and Data.define support
Recognize `Foo = Struct.new(:bar, :baz)` and `Foo = Data.define(:x, :y)` patterns and generate appropriate class definitions with: - attr_reader for all members - attr_writer for Struct members (Data is frozen) - initialize with positional args (Struct) or keyword args (Data) - Block body support for additional method definitions Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 2e4e489 commit 51e939b

7 files changed

Lines changed: 258 additions & 3 deletions

File tree

lib/typeprof/core/ast.rb

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,13 @@ def self.create_node(raw_node, lenv, use_result = true, allow_meta = false)
102102
when :constant_read_node, :constant_path_node
103103
ConstantReadNode.new(raw_node, lenv)
104104
when :constant_write_node, :constant_path_write_node
105-
ConstantWriteNode.new(raw_node, AST.create_node(raw_node.value, lenv), lenv)
105+
if (members = detect_struct_new(raw_node.value))
106+
StructNewNode.new(raw_node, members, :struct, lenv)
107+
elsif (members = detect_data_define(raw_node.value))
108+
StructNewNode.new(raw_node, members, :data, lenv)
109+
else
110+
ConstantWriteNode.new(raw_node, AST.create_node(raw_node.value, lenv), lenv)
111+
end
106112
when :constant_operator_write_node
107113
read = ConstantReadNode.new(raw_node, lenv)
108114
rhs = OperatorNode.new(raw_node, read, lenv)
@@ -526,5 +532,31 @@ def self.create_rbs_type(raw_decl, lenv)
526532
raise "unknown RBS type: #{ raw_decl.class }"
527533
end
528534
end
535+
536+
def self.detect_struct_new(raw_value)
537+
return nil unless raw_value.type == :call_node
538+
return nil unless raw_value.name == :new
539+
recv = raw_value.receiver
540+
return nil unless recv&.type == :constant_read_node && recv.name == :Struct
541+
extract_symbol_args(raw_value)
542+
end
543+
544+
def self.detect_data_define(raw_value)
545+
return nil unless raw_value.type == :call_node
546+
return nil unless raw_value.name == :define
547+
recv = raw_value.receiver
548+
return nil unless recv&.type == :constant_read_node && recv.name == :Data
549+
extract_symbol_args(raw_value)
550+
end
551+
552+
def self.extract_symbol_args(raw_call)
553+
return nil unless raw_call.arguments
554+
members = []
555+
raw_call.arguments.arguments.each do |arg|
556+
return nil unless arg.type == :symbol_node
557+
members << arg.value.to_sym
558+
end
559+
members
560+
end
529561
end
530562
end

lib/typeprof/core/ast/meta.rb

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,5 +208,135 @@ def install0(genv)
208208
Source.new
209209
end
210210
end
211+
212+
class StructNewNode < Node
213+
def initialize(raw_node, members, kind, lenv)
214+
super(raw_node, lenv)
215+
case raw_node.type
216+
when :constant_write_node
217+
@static_cpath = lenv.cref.cpath + [raw_node.name]
218+
when :constant_path_write_node
219+
@static_cpath = AST.parse_cpath(raw_node.target, lenv.cref)
220+
else
221+
raise
222+
end
223+
@members = members
224+
@kind = kind # :struct or :data
225+
226+
# Parse block body if present (Struct.new(:foo) do ... end)
227+
raw_value = raw_node.value
228+
if raw_value.block && raw_value.block.type == :block_node && raw_value.block.body
229+
ncref = CRef.new(@static_cpath, :instance, nil, lenv.cref)
230+
nlenv = LocalEnv.new(lenv.file_context, ncref, {}, [])
231+
@block_body = AST.create_node(raw_value.block.body, nlenv)
232+
end
233+
end
234+
235+
attr_reader :static_cpath, :members, :kind, :block_body
236+
237+
def subnodes = { block_body: }
238+
def attrs = { static_cpath:, members:, kind: }
239+
240+
# Interface expected by MethodDefBox
241+
def req_positionals = @kind == :struct ? @members : []
242+
def opt_positionals = []
243+
def rest_positionals = nil
244+
def post_positionals = []
245+
def req_keywords = @kind == :data ? @members : []
246+
def opt_keywords = []
247+
def rest_keywords = nil
248+
def no_keywords = @kind == :struct
249+
250+
def define0(genv)
251+
mod = genv.resolve_cpath(@static_cpath)
252+
# add_module_def internally calls get_const(name).add_def(self)
253+
cdef = mod.add_module_def(genv, self)
254+
@members.each do |member|
255+
ive = genv.resolve_ivar(@static_cpath, false, member)
256+
ive.add_def(self)
257+
end
258+
@block_body.define(genv) if @block_body
259+
cdef
260+
end
261+
262+
def define_copy(genv)
263+
mod = genv.resolve_cpath(@static_cpath)
264+
mod.add_module_def(genv, self)
265+
mod.remove_module_def(genv, @prev_node)
266+
@members.each do |member|
267+
ive = genv.resolve_ivar(@static_cpath, false, member)
268+
ive.add_def(self)
269+
ive.remove_def(@prev_node)
270+
end
271+
super(genv)
272+
end
273+
274+
def undefine0(genv)
275+
mod = genv.resolve_cpath(@static_cpath)
276+
mod.remove_module_def(genv, self)
277+
@members.each do |member|
278+
ive = genv.resolve_ivar(@static_cpath, false, member)
279+
ive.remove_def(self)
280+
end
281+
@block_body.undefine(genv) if @block_body
282+
end
283+
284+
def install0(genv)
285+
# Register the class singleton type as the constant value
286+
mod_val = Source.new(Type::Singleton.new(genv, genv.resolve_cpath(@static_cpath)))
287+
if @static_cpath
288+
@changes.add_edge(genv, mod_val, @static_ret.vtx)
289+
end
290+
291+
cpath = @static_cpath
292+
@members.each do |member|
293+
# Use bare `:member` (not `:@member`) so the slot can't collide with a
294+
# user-written @member ivar — Struct/Data fields are not real ivars.
295+
ivar_box = @changes.add_ivar_read_box(genv, cpath, false, member)
296+
ret_box = @changes.add_escape_box(genv, ivar_box.ret)
297+
@changes.add_method_def_box(genv, cpath, false, member, FormalArguments::Empty, [ret_box])
298+
299+
if @kind == :struct
300+
# attr_writer (Struct only, Data is frozen)
301+
ive = genv.resolve_ivar(cpath, false, member)
302+
vtx = Vertex.new(self)
303+
@changes.add_edge(genv, vtx, ive.vtx)
304+
writer_ret = @changes.add_escape_box(genv, vtx)
305+
f_args = FormalArguments.new([vtx], [], nil, [], [], [], nil, nil)
306+
@changes.add_method_def_box(genv, cpath, false, :"#{ member }=", f_args, [writer_ret])
307+
end
308+
end
309+
310+
# initialize
311+
init_vtxs = @members.map do |member|
312+
ive = genv.resolve_ivar(cpath, false, member)
313+
vtx = Vertex.new(self)
314+
@changes.add_edge(genv, vtx, ive.vtx)
315+
vtx
316+
end
317+
init_ret = @changes.add_escape_box(genv, Source.new(genv.nil_type))
318+
if @kind == :struct
319+
init_f_args = FormalArguments.new(init_vtxs, [], nil, [], [], [], nil, nil)
320+
else
321+
# Data.define uses keyword arguments
322+
init_f_args = FormalArguments.new([], [], nil, [], init_vtxs, [], nil, nil)
323+
end
324+
@changes.add_method_def_box(genv, cpath, false, :initialize, init_f_args, [init_ret])
325+
326+
# Struct.[] is an alias for Struct.new
327+
if @kind == :struct
328+
self_ret = @changes.add_escape_box(genv, Source.new(Type::Instance.new(genv, genv.resolve_cpath(cpath), [])))
329+
@changes.add_method_def_box(genv, cpath, true, :[], init_f_args, [self_ret])
330+
end
331+
332+
# Install block body (additional method definitions)
333+
if @block_body
334+
@block_body.lenv.locals[:"*self"] = @block_body.lenv.cref.get_self(genv)
335+
@block_body.install(genv)
336+
end
337+
338+
mod_val
339+
end
340+
end
211341
end
212342
end

lib/typeprof/core/env/module_entity.rb

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -261,6 +261,8 @@ def find_superclass_const_read
261261
next
262262
when AST::ModuleNode
263263
return nil
264+
when AST::StructNewNode
265+
return [] # inherits from Object (Struct < Object)
264266
else
265267
raise
266268
end

lib/typeprof/core/graph/box.rb

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -935,7 +935,9 @@ def show(output_parameter_names)
935935
@f_args.post_positionals.each do |var|
936936
args << Type.strip_parens(var.show)
937937
end
938-
if @node.is_a?(AST::DefNode)
938+
if @node.respond_to?(:req_keywords) &&
939+
@node.req_keywords.size == @f_args.req_keywords.size &&
940+
@node.opt_keywords.size == @f_args.opt_keywords.size
939941
@node.req_keywords.zip(@f_args.req_keywords) do |name, f_vtx|
940942
args << "#{ name }: #{Type.strip_parens(f_vtx.show)}"
941943
end

lib/typeprof/core/service.rb

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -472,7 +472,7 @@ def dump_declarations(path)
472472
out << " " * stack.size + "end"
473473
end
474474
end
475-
when AST::ClassNode, AST::SingletonClassNode
475+
when AST::ClassNode, AST::SingletonClassNode, AST::StructNewNode
476476
if node.static_cpath
477477
next if stack.any? { node.is_a?(AST::SingletonClassNode) && (_1.is_a?(AST::ClassNode) || _1.is_a?(AST::ModuleNode)) && node.static_cpath == _1.static_cpath }
478478

@@ -496,6 +496,10 @@ def dump_declarations(path)
496496
out << " " * stack.size + "include #{ inc_mod.show_cpath }"
497497
end
498498
end
499+
# Output method definitions from meta nodes (StructNewNode etc.)
500+
node.boxes(:mdef) do |mdef|
501+
out << " " * stack.size + "def #{ mdef.singleton ? "self." : "" }#{ mdef.mid }: " + mdef.show(@options[:output_parameter_names])
502+
end
499503
else
500504
stack.pop
501505
out << " " * stack.size + "end"
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
## update
2+
class Point < Struct.new(:x, :y)
3+
end
4+
Point.new(1, "hello").x
5+
6+
## assert
7+
class Point < Struct[untyped]
8+
def x: -> Integer
9+
def y: -> String
10+
def x=: (Integer) -> Integer
11+
def y=: (untyped) -> untyped
12+
def initialize: (Integer, String) -> void
13+
def self.[]: (Integer, String) -> Point
14+
end

scenario/misc/struct_new.rb

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
## update
2+
Foo = Struct.new(:bar, :baz)
3+
f = Foo.new(1, "hello")
4+
f.bar
5+
f.baz
6+
f.bar = 2
7+
g = Foo[3, "world"]
8+
9+
## assert
10+
class Foo
11+
def bar: -> Integer
12+
def bar=: (Integer) -> Integer
13+
def baz: -> String
14+
def baz=: (untyped) -> untyped
15+
def initialize: (Integer, String) -> void
16+
def self.[]: (Integer, String) -> Foo
17+
end
18+
19+
## update
20+
Pt = Data.define(:x, :y)
21+
p = Pt.new(x: 1, y: "hello")
22+
p.x
23+
p.y
24+
25+
## assert
26+
class Pt
27+
def x: -> Integer
28+
def y: -> String
29+
def initialize: (x: Integer, y: String) -> void
30+
end
31+
32+
## update
33+
Bar = Struct.new(:n) do
34+
def double
35+
n * 2
36+
end
37+
end
38+
Bar.new(5).double
39+
40+
## assert
41+
class Bar
42+
def n: -> Integer
43+
def n=: (untyped) -> untyped
44+
def initialize: (Integer) -> void
45+
def self.[]: (Integer) -> Bar
46+
def double: -> Integer
47+
end
48+
49+
## update
50+
# The Struct member `v` is not a real Ruby ivar, so a user-written @v inside
51+
# the block body must not share the member's type.
52+
Baz = Struct.new(:v) do
53+
def set_label
54+
@v = "label"
55+
end
56+
def ivar
57+
@v
58+
end
59+
end
60+
Baz.new(42).v
61+
Baz.new(42).ivar
62+
63+
## assert
64+
class Baz
65+
def v: -> Integer
66+
def v=: (untyped) -> untyped
67+
def initialize: (Integer) -> void
68+
def self.[]: (Integer) -> Baz
69+
def set_label: -> String
70+
def ivar: -> String
71+
end

0 commit comments

Comments
 (0)