Skip to content

Commit 3900495

Browse files
mediremimediremi
committed
Add type expansion refactor from rescript-lang#7566 
1 parent 722896c commit 3900495

File tree

1 file changed

+253
-0
lines changed

1 file changed

+253
-0
lines changed

‎analysis/src/TypeUtils.ml

Lines changed: 253 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1291,3 +1291,256 @@ let completionPathFromMaybeBuiltin path =
12911291
(* Route Stdlib_X to Stdlib.X for proper completions without the Stdlib_ prefix *)
12921292
Some (String.split_on_char '_' mainModule)
12931293
| _ -> None)
1294+
1295+
module ExpandType = struct
1296+
type expandTypeInput =
1297+
| TypeExpr of {
1298+
typeExpr: Types.type_expr;
1299+
name: string Location.loc option;
1300+
env: QueryEnv.t;
1301+
}
1302+
| TypeDecl of {
1303+
typeDecl: Types.type_declaration;
1304+
name: string Location.loc;
1305+
env: QueryEnv.t;
1306+
}
1307+
1308+
type expandTypeReturn = {
1309+
mainTypes: expandTypeInput list;
1310+
relatedTypes: expandTypeInput list;
1311+
}
1312+
1313+
module TypeIdSet = Set.Make (String)
1314+
1315+
let expandTypes (input : expandTypeInput) ~(full : SharedTypes.full) =
1316+
let rootEnv = QueryEnv.fromFile full.file in
1317+
1318+
let expandTypeInputToKey = function
1319+
| TypeExpr {name; env} ->
1320+
typeId ~env
1321+
~name:
1322+
(match name with
1323+
| None -> Location.mkloc "<unknown>" Location.none
1324+
| Some n -> n)
1325+
| TypeDecl {name; env} -> typeId ~env ~name
1326+
in
1327+
1328+
let deduplicateAndRemoveAlreadyPresent mainTypes relatedTypes =
1329+
let mainIds = ref TypeIdSet.empty in
1330+
let dedupedMain =
1331+
mainTypes
1332+
|> List.fold_left
1333+
(fun acc item ->
1334+
let id = expandTypeInputToKey item in
1335+
if TypeIdSet.mem id !mainIds then acc
1336+
else (
1337+
mainIds := TypeIdSet.add id !mainIds;
1338+
item :: acc))
1339+
[]
1340+
|> List.rev
1341+
in
1342+
1343+
let relatedIds = ref TypeIdSet.empty in
1344+
let dedupedRelated =
1345+
relatedTypes
1346+
|> List.fold_left
1347+
(fun acc item ->
1348+
let id = expandTypeInputToKey item in
1349+
if TypeIdSet.mem id !mainIds || TypeIdSet.mem id !relatedIds then
1350+
acc
1351+
else (
1352+
relatedIds := TypeIdSet.add id !relatedIds;
1353+
item :: acc))
1354+
[]
1355+
|> List.rev
1356+
in
1357+
1358+
(dedupedMain, dedupedRelated)
1359+
in
1360+
1361+
let rec followTypeAliases acc (typeExpr : Types.type_expr) =
1362+
match typeExpr.desc with
1363+
| Tlink t1 | Tsubst t1 | Tpoly (t1, []) -> followTypeAliases acc t1
1364+
| Tconstr (path, typeArgs, _) -> (
1365+
match
1366+
References.digConstructor ~env:rootEnv ~package:full.package path
1367+
with
1368+
| Some
1369+
( env,
1370+
{
1371+
name;
1372+
item = {decl = {type_manifest = Some t1; type_params} as decl};
1373+
} ) ->
1374+
let instantiated =
1375+
instantiateType ~typeParams:type_params ~typeArgs t1
1376+
in
1377+
let currentAlias = TypeDecl {typeDecl = decl; name; env} in
1378+
followTypeAliases (currentAlias :: acc) instantiated
1379+
| Some (env, {name; item = {decl}}) ->
1380+
TypeDecl {typeDecl = decl; name; env} :: acc
1381+
| None -> acc)
1382+
| _ -> acc
1383+
in
1384+
1385+
let rec findFinalConcreteType (typeExpr : Types.type_expr) =
1386+
match typeExpr.desc with
1387+
| Tlink t1 | Tsubst t1 | Tpoly (t1, []) -> findFinalConcreteType t1
1388+
| Tconstr (path, typeArgs, _) -> (
1389+
match
1390+
References.digConstructor ~env:rootEnv ~package:full.package path
1391+
with
1392+
| Some (_env, {item = {decl = {type_manifest = Some t1; type_params}}})
1393+
->
1394+
let instantiated =
1395+
instantiateType ~typeParams:type_params ~typeArgs t1
1396+
in
1397+
findFinalConcreteType instantiated
1398+
| _ -> typeExpr)
1399+
| _ -> typeExpr
1400+
in
1401+
1402+
let rec extractRelevantTypesFromTypeExpr ?(depth = 0)
1403+
(typeExpr : Types.type_expr) =
1404+
if depth > 1 then []
1405+
else
1406+
match typeExpr.desc with
1407+
| Tlink t1 | Tsubst t1 | Tpoly (t1, []) ->
1408+
extractRelevantTypesFromTypeExpr ~depth t1
1409+
| Tconstr (path, typeArgs, _) ->
1410+
let constructorTypes =
1411+
match
1412+
References.digConstructor ~env:rootEnv ~package:full.package path
1413+
with
1414+
| Some (env, {name; item = {kind = Record fields; decl}}) ->
1415+
TypeDecl {typeDecl = decl; name; env}
1416+
::
1417+
(if depth = 0 then
1418+
fields
1419+
|> List.fold_left
1420+
(fun acc field ->
1421+
acc
1422+
@ extractRelevantTypesFromTypeExpr ~depth:(depth + 1)
1423+
field.typ)
1424+
[]
1425+
else [])
1426+
| Some (env, {name; item = {kind = Variant constructors; decl}}) ->
1427+
TypeDecl {typeDecl = decl; name; env}
1428+
::
1429+
(if depth = 0 then
1430+
constructors
1431+
|> List.fold_left
1432+
(fun acc (constructor : Constructor.t) ->
1433+
match constructor.args with
1434+
| Args args ->
1435+
args
1436+
|> List.fold_left
1437+
(fun acc (argType, _) ->
1438+
acc
1439+
@ extractRelevantTypesFromTypeExpr
1440+
~depth:(depth + 1) argType)
1441+
acc
1442+
| InlineRecord fields ->
1443+
fields
1444+
|> List.fold_left
1445+
(fun acc field ->
1446+
acc
1447+
@ extractRelevantTypesFromTypeExpr
1448+
~depth:(depth + 1) field.typ)
1449+
acc)
1450+
[]
1451+
else [])
1452+
| Some (_env, {item = {decl = {type_manifest = Some t1}}}) ->
1453+
extractRelevantTypesFromTypeExpr ~depth t1
1454+
| _ -> []
1455+
in
1456+
let typeArgTypes =
1457+
typeArgs
1458+
|> List.fold_left
1459+
(fun acc typeArg ->
1460+
acc
1461+
@ extractRelevantTypesFromTypeExpr ~depth:(depth + 1) typeArg)
1462+
[]
1463+
in
1464+
constructorTypes @ typeArgTypes
1465+
| Tvariant {row_fields} when depth = 0 ->
1466+
row_fields
1467+
|> List.fold_left
1468+
(fun acc (_label, field) ->
1469+
match field with
1470+
| Types.Rpresent (Some typeExpr) ->
1471+
acc
1472+
@ extractRelevantTypesFromTypeExpr ~depth:(depth + 1)
1473+
typeExpr
1474+
| Reither (_, typeExprs, _, _) ->
1475+
typeExprs
1476+
|> List.fold_left
1477+
(fun acc typeExpr ->
1478+
acc
1479+
@ extractRelevantTypesFromTypeExpr ~depth:(depth + 1)
1480+
typeExpr)
1481+
acc
1482+
| _ -> acc)
1483+
[]
1484+
| _ -> []
1485+
in
1486+
1487+
let extractRelevantTypesFromTypeDecl (typeDecl : Types.type_declaration) =
1488+
match typeDecl.type_manifest with
1489+
| Some typeExpr -> extractRelevantTypesFromTypeExpr typeExpr
1490+
| None -> (
1491+
match typeDecl.type_kind with
1492+
| Type_record (label_declarations, _) ->
1493+
label_declarations
1494+
|> List.fold_left
1495+
(fun acc (label_decl : Types.label_declaration) ->
1496+
acc
1497+
@ extractRelevantTypesFromTypeExpr ~depth:1 label_decl.ld_type)
1498+
[]
1499+
| Type_variant constructor_declarations ->
1500+
constructor_declarations
1501+
|> List.fold_left
1502+
(fun acc (constructor_decl : Types.constructor_declaration) ->
1503+
match constructor_decl.cd_args with
1504+
| Cstr_tuple type_exprs ->
1505+
type_exprs
1506+
|> List.fold_left
1507+
(fun acc type_expr ->
1508+
acc
1509+
@ extractRelevantTypesFromTypeExpr ~depth:1 type_expr)
1510+
acc
1511+
| Cstr_record label_declarations ->
1512+
label_declarations
1513+
|> List.fold_left
1514+
(fun acc (label_decl : Types.label_declaration) ->
1515+
acc
1516+
@ extractRelevantTypesFromTypeExpr ~depth:1
1517+
label_decl.ld_type)
1518+
acc)
1519+
[]
1520+
| Type_abstract | Type_open -> [])
1521+
in
1522+
1523+
match input with
1524+
| TypeExpr {typeExpr; name; env} ->
1525+
let aliases = followTypeAliases [] typeExpr in
1526+
let mainTypesRaw = TypeExpr {typeExpr; name; env} :: aliases in
1527+
1528+
(* Extract related types from the final concrete type *)
1529+
let finalConcreteType = findFinalConcreteType typeExpr in
1530+
let relatedTypesRaw =
1531+
extractRelevantTypesFromTypeExpr finalConcreteType
1532+
in
1533+
1534+
let mainTypes, relatedTypes =
1535+
deduplicateAndRemoveAlreadyPresent mainTypesRaw relatedTypesRaw
1536+
in
1537+
{mainTypes; relatedTypes}
1538+
| TypeDecl {typeDecl} ->
1539+
let mainTypes = [input] in
1540+
let relatedTypesRaw = extractRelevantTypesFromTypeDecl typeDecl in
1541+
1542+
let _, relatedTypes =
1543+
deduplicateAndRemoveAlreadyPresent mainTypes relatedTypesRaw
1544+
in
1545+
{mainTypes; relatedTypes}
1546+
end

0 commit comments

Comments
 (0)