xref: /llvm-project/mlir/test/python/ir/location.py (revision a77250fd782530f42a90f8562bcef0eb26abb010) 
1# RUN: %PYTHON %s | FileCheck %s
2
3import gc
4from mlir.ir import *
5
6
7def run(f):
8    print("\nTEST:", f.__name__)
9    f()
10    gc.collect()
11    assert Context._get_live_count() == 0
12
13
14# CHECK-LABEL: TEST: testUnknown
15def testUnknown():
16    with Context() as ctx:
17        loc = Location.unknown()
18    assert loc.context is ctx
19    ctx = None
20    gc.collect()
21    # CHECK: unknown str: loc(unknown)
22    print("unknown str:", str(loc))
23    # CHECK: unknown repr: loc(unknown)
24    print("unknown repr:", repr(loc))
25
26
27run(testUnknown)
28
29
30# CHECK-LABEL: TEST: testLocationAttr
31def testLocationAttr():
32    with Context() as ctxt:
33        loc = Location.unknown()
34        attr = loc.attr
35        clone = Location.from_attr(attr)
36    gc.collect()
37    # CHECK: loc: loc(unknown)
38    print("loc:", str(loc))
39    # CHECK: clone: loc(unknown)
40    print("clone:", str(clone))
41    assert loc == clone
42
43
44run(testLocationAttr)
45
46# CHECK-LABEL: TEST: testFileLineCol
47def testFileLineCol():
48    with Context() as ctx:
49        loc = Location.file("foo.txt", 123, 56)
50        range = Location.file("foo.txt", 123, 56, 123, 100)
51    ctx = None
52    gc.collect()
53    # CHECK: file str: loc("foo.txt":123:56)
54    print("file str:", str(loc))
55    # CHECK: file repr: loc("foo.txt":123:56)
56    print("file repr:", repr(loc))
57    # CHECK: file range str: loc("foo.txt":123:56 to :100)
58    print("file range str:", str(range))
59    # CHECK: file range repr: loc("foo.txt":123:56 to :100)
60    print("file range repr:", repr(range))
61
62
63run(testFileLineCol)
64
65
66# CHECK-LABEL: TEST: testName
67def testName():
68    with Context() as ctx:
69        loc = Location.name("nombre")
70        locWithChildLoc = Location.name("naam", loc)
71    ctx = None
72    gc.collect()
73    # CHECK: file str: loc("nombre")
74    print("file str:", str(loc))
75    # CHECK: file repr: loc("nombre")
76    print("file repr:", repr(loc))
77    # CHECK: file str: loc("naam"("nombre"))
78    print("file str:", str(locWithChildLoc))
79    # CHECK: file repr: loc("naam"("nombre"))
80    print("file repr:", repr(locWithChildLoc))
81
82
83run(testName)
84
85
86# CHECK-LABEL: TEST: testCallSite
87def testCallSite():
88    with Context() as ctx:
89        loc = Location.callsite(
90            Location.file("foo.text", 123, 45),
91            [Location.file("util.foo", 379, 21), Location.file("main.foo", 100, 63)],
92        )
93    ctx = None
94    # CHECK: file str: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
95    print("file str:", str(loc))
96    # CHECK: file repr: loc(callsite("foo.text":123:45 at callsite("util.foo":379:21 at "main.foo":100:63))
97    print("file repr:", repr(loc))
98
99
100run(testCallSite)
101
102
103# CHECK-LABEL: TEST: testFused
104def testFused():
105    with Context() as ctx:
106        loc_single = Location.fused([Location.name("apple")])
107        loc = Location.fused([Location.name("apple"), Location.name("banana")])
108        attr = Attribute.parse('"sauteed"')
109        loc_attr = Location.fused(
110            [Location.name("carrot"), Location.name("potatoes")], attr
111        )
112        loc_empty = Location.fused([])
113        loc_empty_attr = Location.fused([], attr)
114        loc_single_attr = Location.fused([Location.name("apple")], attr)
115    ctx = None
116    # CHECK: file str: loc("apple")
117    print("file str:", str(loc_single))
118    # CHECK: file repr: loc("apple")
119    print("file repr:", repr(loc_single))
120    # CHECK: file str: loc(fused["apple", "banana"])
121    print("file str:", str(loc))
122    # CHECK: file repr: loc(fused["apple", "banana"])
123    print("file repr:", repr(loc))
124    # CHECK: file str: loc(fused<"sauteed">["carrot", "potatoes"])
125    print("file str:", str(loc_attr))
126    # CHECK: file repr: loc(fused<"sauteed">["carrot", "potatoes"])
127    print("file repr:", repr(loc_attr))
128    # CHECK: file str: loc(unknown)
129    print("file str:", str(loc_empty))
130    # CHECK: file repr: loc(unknown)
131    print("file repr:", repr(loc_empty))
132    # CHECK: file str: loc(fused<"sauteed">[unknown])
133    print("file str:", str(loc_empty_attr))
134    # CHECK: file repr: loc(fused<"sauteed">[unknown])
135    print("file repr:", repr(loc_empty_attr))
136    # CHECK: file str: loc(fused<"sauteed">["apple"])
137    print("file str:", str(loc_single_attr))
138    # CHECK: file repr: loc(fused<"sauteed">["apple"])
139    print("file repr:", repr(loc_single_attr))
140
141
142run(testFused)
143
144
145# CHECK-LABEL: TEST: testLocationCapsule
146def testLocationCapsule():
147    with Context() as ctx:
148        loc1 = Location.file("foo.txt", 123, 56)
149    # CHECK: mlir.ir.Location._CAPIPtr
150    loc_capsule = loc1._CAPIPtr
151    print(loc_capsule)
152    loc2 = Location._CAPICreate(loc_capsule)
153    assert loc2 == loc1
154    assert loc2.context is ctx
155
156
157run(testLocationCapsule)
158