Fix %install issues with modulemaps (#58)

Author: vvmnnnkv

register.py

@@ -56,6 +56,7 @@ def make_kernel_env(args):
             kernel_env['LD_LIBRARY_PATH'] = '%s/usr/lib/swift/linux' % args.swift_toolchain
             kernel_env['REPL_SWIFT_PATH'] = '%s/usr/bin/repl_swift' % args.swift_toolchain
             kernel_env['SWIFT_BUILD_PATH'] = '%s/usr/bin/swift-build' % args.swift_toolchain
+            kernel_env['SWIFT_PACKAGE_PATH'] = '%s/usr/bin/swift-package' % args.swift_toolchain
         elif platform.system() == 'Darwin':
             kernel_env['PYTHONPATH'] = '%s/System/Library/PrivateFrameworks/LLDB.framework/Resources/Python' % args.swift_toolchain
             kernel_env['LD_LIBRARY_PATH'] = '%s/usr/lib/swift/macosx' % args.swift_toolchain
@@ -118,6 +119,10 @@ def validate_kernel_env(kernel_env):
             not os.path.isfile(kernel_env['SWIFT_BUILD_PATH']):
         raise Exception('swift-build binary not found at %s' %
                         kernel_env['SWIFT_BUILD_PATH'])
+    if 'SWIFT_PACKAGE_PATH' in kernel_env and \
+            not os.path.isfile(kernel_env['SWIFT_PACKAGE_PATH']):
+        raise Exception('swift-package binary not found at %s' %
+                        kernel_env['SWIFT_PACKAGE_PATH'])
     if 'PYTHON_LIBRARY' in kernel_env and \
             not os.path.isfile(kernel_env['PYTHON_LIBRARY']):
         raise Exception('python library not found at %s' %

swift_kernel.py

@@ -30,6 +30,8 @@
 import textwrap
 import time
 import threading
+import sqlite3
+import json
 
 from ipykernel.kernelbase import Kernel
 from jupyter_client.jsonutil import squash_dates
@@ -494,6 +496,12 @@ def _install_packages(self, packages, swiftpm_flags):
                     'Install Error: Cannot install packages because '
                     'SWIFT_BUILD_PATH is not specified.')
 
+        swift_package_path = os.environ.get('SWIFT_PACKAGE_PATH')
+        if swift_package_path is None:
+            raise PackageInstallException(
+                    'Install Error: Cannot install packages because '
+                    'SWIFT_PACKAGE_PATH is not specified.')
+
         swift_import_search_path = os.environ.get('SWIFT_IMPORT_SEARCH_PATH')
         if swift_import_search_path is None:
             raise PackageInstallException(
@@ -591,27 +599,74 @@ def _install_packages(self, packages, swiftpm_flags):
 
         # == Copy .swiftmodule and modulemap files to SWIFT_IMPORT_SEARCH_PATH ==
 
-        for filename in glob.glob(os.path.join(bin_dir, '*.swiftmodule')):
+        build_db_file = os.path.join(package_base_path, '.build', 'build.db')
+        if not os.path.exists(build_db_file):
+            raise PackageInstallException('build.db is missing')
+
+        # Execute swift-package show-dependencies to get all dependencies' paths
+        dependencies_result = subprocess.run(
+            [swift_package_path, 'show-dependencies', '--format', 'json'],
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+            cwd=package_base_path)
+        dependencies_json = dependencies_result.stdout.decode('utf8')
+        dependencies_obj = json.loads(dependencies_json)
+
+        def flatten_deps_paths(dep):
+            paths = []
+            paths.append(dep["path"])
+            if dep["dependencies"]:
+                for d in dep["dependencies"]:
+                    paths.extend(flatten_deps_paths(d))
+            return paths
+
+        # Make list of paths where we expect .swiftmodule and .modulemap files of dependencies
+        dependencies_paths = [package_base_path]
+        dependencies_paths = flatten_deps_paths(dependencies_obj)
+        dependencies_paths = list(set(dependencies_paths))
+
+        def is_valid_dependency(path):
+            for p in dependencies_paths:
+                if path.startswith(p): return True
+            return False
+
+        # Query to get build files list from build.db
+        # SUBSTR because string starts with "N" (why?)
+        SQL_FILES_SELECT = "SELECT SUBSTR(key, 2) FROM 'key_names' WHERE key LIKE ?"
+
+        # Connect to build.db
+        db_connection = sqlite3.connect(build_db_file)
+        cursor = db_connection.cursor()
+
+        # Process *.swiftmodules files
+        cursor.execute(SQL_FILES_SELECT, ['%.swiftmodule'])
+        swift_modules = [row[0] for row in cursor.fetchall() if is_valid_dependency(row[0])]
+        for filename in swift_modules:
             shutil.copy(filename, swift_import_search_path)
 
-        # The modulemap files appear in a few different places. Add all of
-        # them.
-        # TODO: Figure out if there is a principled way to figure out where
-        # all the modulemap files are.
-        modulemap_files = glob.glob(
-                os.path.join(bin_dir, '**/module.modulemap'), recursive=True)
-        modulemap_files += glob.glob(
-                os.path.join(package_base_path,
-                             '.build/checkouts/**/module.modulemap'),
-                recursive=True)
+        # Process modulemap files
+        cursor.execute(SQL_FILES_SELECT, ['%/module.modulemap'])
+        modulemap_files = [row[0] for row in cursor.fetchall() if is_valid_dependency(row[0])]
         for index, filename in enumerate(modulemap_files):
             # Create a separate directory for each modulemap file because the
             # ClangImporter requires that they are all named
             # "module.modulemap".
             modulemap_dest = os.path.join(swift_import_search_path,
                                           'modulemap-%d' % index)
             os.makedirs(modulemap_dest, exist_ok=True)
-            shutil.copy(filename, modulemap_dest)
+            src_folder, src_filename = os.path.split(filename)
+            dst_path = os.path.join(modulemap_dest, src_filename)
+            # Make all relative header paths in module.modulemap absolute
+            # because we copy file to different location
+            with open(filename, encoding='utf8') as file:
+                modulemap_contents = file.read()
+                modulemap_contents = re.sub(
+                    r'header\s+"(.*?)"',
+                    lambda m: 'header "%s"' %
+                        (m.group(1) if os.path.isabs(m.group(1)) else os.path.abspath(os.path.join(src_folder, m.group(1)))),
+                    modulemap_contents
+                )
+                with open(dst_path, 'w', encoding='utf8') as outfile:
+                    outfile.write(modulemap_contents)
 
         # == dlopen the shared lib ==