Source code for ffcx.formatting

# Copyright (C) 2009-2018 Anders Logg and Garth N. Wells
#
# This file is part of FFCx. (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Compiler stage 5: Code formatting.

This module implements the formatting of UFC code from a given
dictionary of generated C++ code for the body of each UFC function.

It relies on templates for UFC code available as part of the module
ufcx_utils.

"""

import logging
import os
import pprint
import textwrap

from ffcx import __version__ as FFCX_VERSION
from ffcx.codegeneration import __version__ as UFC_VERSION

logger = logging.getLogger("ffcx")

FORMAT_TEMPLATE = {
    "ufc comment":
    """\
// This code conforms with the UFC specification version {ufcx_version}
// and was automatically generated by FFCx version {ffcx_version}.
""",
    "dolfinx comment":
    """\
// This code conforms with the UFC specification version {ufcx_version}
// and was automatically generated by FFCx version {ffcx_version}.
//
""",
    "header_h":
    """
#pragma once

""",
    "header_c":
    """
""",
}

c_extern_pre = """
#ifdef __cplusplus
extern "C" {
#endif
"""

c_extern_post = """
#ifdef __cplusplus
}
#endif
"""


def format_code(code, options: dict):
    """Format given code in UFC format. Returns two strings with header and source file contents."""
    logger.info(79 * "*")
    logger.info("Compiler stage 5: Formatting code")
    logger.info(79 * "*")

    # Generate code for comment at top of file
    code_h_pre = _generate_comment(options) + "\n"
    code_c_pre = _generate_comment(options) + "\n"

    # Generate code for header
    code_h_pre += FORMAT_TEMPLATE["header_h"]
    code_c_pre += FORMAT_TEMPLATE["header_c"]

    # Generate includes and add to preamble
    includes_h, includes_c = _generate_includes(options)
    code_h_pre += includes_h
    code_c_pre += includes_c

    # Enclose header with 'extern "C"'
    code_h_pre += c_extern_pre
    code_h_post = c_extern_post

    code_h = ""
    code_c = ""

    for parts_code in code:
        code_h += "".join([c[0] for c in parts_code])
        code_c += "".join([c[1] for c in parts_code])

    # Add headers to body
    code_h = code_h_pre + code_h + code_h_post
    code_c = code_c_pre + code_c

    return code_h, code_c


def write_code(code_h, code_c, prefix, output_dir):
    _write_file(code_h, prefix, ".h", output_dir)
    _write_file(code_c, prefix, ".c", output_dir)


def _write_file(output, prefix, postfix, output_dir):
    """Write generated code to file."""
    filename = os.path.join(output_dir, prefix + postfix)
    with open(filename, "w") as hfile:
        hfile.write(output)


def _generate_comment(options):
    """Generate code for comment on top of file."""
    # Generate top level comment
    comment = FORMAT_TEMPLATE["ufc comment"].format(ffcx_version=FFCX_VERSION, ufcx_version=UFC_VERSION)

    # Add option information
    comment += "//\n"
    comment += "// This code was generated with the following options:\n"
    comment += "//\n"
    comment += textwrap.indent(pprint.pformat(options), "//  ")
    comment += "\n"

    return comment


def _generate_includes(options: dict):

    default_h_includes = [
        "#include <ufcx.h>",
    ]

    default_c_includes = [
        "#include <math.h>",  # This should really be set by the backend
        "#include <stdalign.h>",  # This should really be set by the backend
        "#include <stdlib.h>",  # This should really be set by the backend
        "#include <string.h>",  # This should really be set by the backend
        "#include <ufcx.h>"
    ]

    if "_Complex" in options["scalar_type"]:
        default_c_includes += ["#include <complex.h>"]

    s_h = set(default_h_includes)
    s_c = set(default_c_includes)

    includes_h = "\n".join(sorted(s_h)) + "\n" if s_h else ""
    includes_c = "\n".join(sorted(s_c)) + "\n" if s_c else ""

    return includes_h, includes_c