# Copyright (c) Meta Platforms, Inc. and affiliates.

schema python.branches.1 {
import builtin.1
import python.4
import src

predicate Root: string

predicate FileToRoot:
  {
    file: src.File,
    root: Root,
  }

predicate CommonDependency: Root

predicate TargetLocation:
  {
    file: src.File,
    target: python.Name,
    decl_location: python.DeclarationLocation
  }
  { File, TargetName, Target } where
  python.DeclarationWithName { TargetName, Decl };
  Target = python.DeclarationLocation { declaration = Decl, file = SourceFile };
  FileToRoot { file = SourceFile, root = Root };
  (FileToRoot { file = File, root = Root } | CommonDependency Root);

predicate DirectXRefsByFile:
  {
    file: src.File,
    target: python.DeclarationLocation,
    source: src.ByteSpan,
  }
  { File, Target, Span } where
  python.XRefsViaNameByFile { File, XRefs };
  { target = TargetName, source = Span } = XRefs[..];
  TargetLocation {File, TargetName, Target};


predicate ImportedDeclaration:
  {
    import_decl: python.DeclarationLocation,
    imported_declaration: python.DeclarationLocation,
  } {ImportDecl, ImportedDeclaration} where
  python.DeclarationLocation { declaration = {imp = { from_name = FromName} }, file = ImportFile} = ImportDecl;
  TargetLocation {ImportFile, FromName, ImportedDeclaration};


predicate ResolveOriginalDeclaration:
  {
    decl_location: python.DeclarationLocation,
    original_decl_location: python.DeclarationLocation,
  }
  {DeclLocation, OriginalDeclaration} where
  OriginalDeclaration = if (
      ImportedDeclaration { DeclLocation, DeclLocation1};
    ) then (
      if (
        ImportedDeclaration { DeclLocation1, DeclLocation2 };
      ) then (
        if (
          ImportedDeclaration { DeclLocation2, DeclLocation3 };
        ) then (
          if (
            ImportedDeclaration { DeclLocation3, DeclLocation4 };
          ) then (
            if (
              ImportedDeclaration { DeclLocation4, DeclLocation5 };
            ) then (
              DeclLocation5 # returns declaration location 5 even if there are other imports
            ) else ( DeclLocation4 )
          ) else ( DeclLocation3 )
        ) else ( DeclLocation2 )
      ) else ( DeclLocation1 )
    ) else ( DeclLocation );


predicate ResolvedXRefsByFile:
  {
      file: src.File,
      target: python.DeclarationLocation,
      source: src.ByteSpan,
  }
  { File, OriginalDeclLocation, Span } where
  DirectXRefsByFile { File, DeclLocation, Span};
  ResolveOriginalDeclaration { DeclLocation, OriginalDeclLocation};
}