[fhir-tools] Add Generation of Extensions to the Health-Tool and Minor Bug Fixes

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/BallerinaPackageGenTool.java

@@ -76,6 +76,7 @@ public TemplateGenerator execute(ToolContext toolContext) throws CodeGenExceptio
             }
             properties.put("packageContext", packageContextGenerator.getPackageContext());
             properties.put("datatypeContext", packageContextGenerator.getPackageContext().getDatatypeTemplateContextMap());
+            properties.put("extensionContext", packageContextGenerator.getPackageContext().getExtensionTemplateContext());
             packageTemplateGenerator.setGeneratorProperties(properties);
             return packageTemplateGenerator;
         }

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/model/ExtensionTemplateContext.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.healthcare.fhir.ballerina.packagegen.tool.model;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.HashMap;
+
+public class ExtensionTemplateContext {
+    private Map<String, DatatypeTemplateContext> extensionDatatypes;
+    private Map<String, Set<String>> extensionSlices;
+    private Map<String, Set<String>> extensionResources;
+
+    public ExtensionTemplateContext() {
+        extensionDatatypes = new TreeMap<>();
+        extensionSlices = new HashMap<>();
+        extensionResources = new HashMap<>();
+    }
+
+    public Map<String, DatatypeTemplateContext> getExtensionDatatypes() {
+        return extensionDatatypes;
+    }
+
+    public void setExtensionDatatypes(Map<String, DatatypeTemplateContext> extensionDatatypes) {
+        this.extensionDatatypes = extensionDatatypes;
+    }
+
+    public Map<String, Set<String>> getExtensionSlices() {
+        return extensionSlices;
+    }
+
+    public void setExtensionSlices(Map<String, Set<String>> extensionSlices) {
+        this.extensionSlices = extensionSlices;
+    }
+
+    public Map<String, Set<String>> getExtensionResources() {
+        return extensionResources;
+    }
+
+    public void setExtensionResources(Map<String, Set<String>> extensionResources) {
+        this.extensionResources = extensionResources;
+    }
+}

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/model/PackageTemplateContext.java

@@ -18,7 +18,6 @@
 
 package org.wso2.healthcare.fhir.ballerina.packagegen.tool.model;
 
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -31,6 +30,7 @@ public class PackageTemplateContext {
     private String internationalPackageName;
     private Map<String, ResourceTemplateContext> resourceTemplateContextMap;
     private Map<String, DatatypeTemplateContext> datatypeTemplateContextMap;
+    private ExtensionTemplateContext extensionTemplateContext;
     private Map<String, String> resourceNameTypeMap;
     private Set<String> dataTypesRegistry;
     private Map<String, String> dependenciesMap;
@@ -80,6 +80,14 @@ public void addDatatypeTemplateContext(String datatypeName, DatatypeTemplateCont
         this.datatypeTemplateContextMap.putIfAbsent(datatypeName, context);
     }
 
+    public ExtensionTemplateContext getExtensionTemplateContext() {
+        return extensionTemplateContext;
+    }
+
+    public void setExtensionTemplateContext(ExtensionTemplateContext extensionTemplateContext) {
+        this.extensionTemplateContext = extensionTemplateContext;
+    }
+
     public Map<String, String> getResourceNameTypeMap() {
         return resourceNameTypeMap;
     }

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/model/ResourceTemplateContext.java

@@ -23,6 +23,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.Map;
 
 /**
  * Class holder for resource related template context
@@ -44,6 +45,7 @@ public class ResourceTemplateContext {
     private HashMap<String, List<Element>> sliceElements = new HashMap<>();
     private HashMap<String, ExtendedElement> extendedElements = new HashMap<>();
     private Set<String> resourceDependencies = new HashSet<>();
+    private Map<String, Set<String>> resourceExtensions = new HashMap<>();
 
     public String getResourceType() {
         return resourceType;
@@ -148,4 +150,12 @@ public Set<String> getResourceDependencies() {
     public void setResourceDependencies(Set<String> resourceDependencies) {
         this.resourceDependencies = resourceDependencies;
     }
+
+    public Map<String, Set<String>> getResourceExtensions() {
+        return resourceExtensions;
+    }
+
+    public void setResourceExtensions(Map<String, Set<String>> resourceExtensions) {
+        this.resourceExtensions = resourceExtensions;
+    }
 }

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/modelgen/AbstractDatatypeContextGenerator.java

@@ -48,19 +48,11 @@ public AbstractDatatypeContextGenerator(FHIRSpecificationData fhirSpecificationD
         populateDatatypeContext();
     }
 
-    public Map<String, FHIRDataTypeDef> getDatatypeDefnMap() {
-        return datatypeDefnMap;
-    }
-
-    public Map<String, DatatypeTemplateContext> getDataTypeTemplateContextMap() {
-        return dataTypeTemplateContextMap;
-    }
-
     protected Map<String, FHIRDataTypeDef> getDataTypeDefnMap() {
         return datatypeDefnMap;
     }
 
-    protected Map<String, DatatypeTemplateContext> datatypeTemplateContextMap() {
+    public Map<String, DatatypeTemplateContext> getDatatypeTemplateContextMap() {
         return dataTypeTemplateContextMap;
     }
 

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/modelgen/AbstractExtensionContextGenerator.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.healthcare.fhir.ballerina.packagegen.tool.modelgen;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.wso2.healthcare.codegen.tool.framework.fhir.core.common.FHIRSpecificationData;
+import org.wso2.healthcare.codegen.tool.framework.fhir.core.model.FHIRDataTypeDef;
+import org.wso2.healthcare.fhir.ballerina.packagegen.tool.DataTypesRegistry;
+import org.wso2.healthcare.fhir.ballerina.packagegen.tool.model.DatatypeTemplateContext;
+import org.wso2.healthcare.fhir.ballerina.packagegen.tool.model.ExtensionTemplateContext;
+
+import java.util.Map;
+import java.util.Set;
+import java.util.HashSet;
+import java.util.HashMap;
+
+public abstract class AbstractExtensionContextGenerator {
+    private static final Log LOG = LogFactory.getLog(AbstractExtensionContextGenerator.class);
+    private final Map<String, FHIRDataTypeDef> extensionDefnMap;
+    private final ExtensionTemplateContext extensionTemplateContext;
+
+    public AbstractExtensionContextGenerator(FHIRSpecificationData fhirSpecificationData) {
+        LOG.info("Started: Extensions Generation");
+
+        this.extensionDefnMap = fhirSpecificationData.getDataTypes();
+        this.extensionTemplateContext = new ExtensionTemplateContext();
+
+        populateExtensionTemplateContext();
+
+        LOG.info("Ended: Extensions Generation");
+    }
+
+    protected Map<String, FHIRDataTypeDef> getExtensionDefnMap() {
+        return extensionDefnMap;
+    }
+
+    public ExtensionTemplateContext getExtensionTemplateContext() {
+        return extensionTemplateContext;
+    }
+
+    private void populateExtensionTemplateContext() {
+        populateBaseExtensionContext();
+        populateSliceExtensionContext();
+        populateExtensionSliceMap();
+    }
+
+    protected abstract void populateBaseExtensionContext();
+
+    protected abstract void populateSliceExtensionContext();
+
+    protected abstract void populateExtensionResourceMap(String identifier, FHIRDataTypeDef extensionDef);
+
+    protected void populateExtensionSliceMap() {
+        Map<String, Set<String>> childExtensionMap = new HashMap<>();
+        Set<String> searchKeys = new HashSet<>();
+        Map<String, DatatypeTemplateContext> extensionDatatypeMap = extensionTemplateContext.getExtensionDatatypes();
+
+        for (Map.Entry<String, DatatypeTemplateContext> contextEntry : extensionDatatypeMap.entrySet()) {
+            if (contextEntry.getKey().contains("http://")) {
+                /// Set Base Extension
+                searchKeys.add(contextEntry.getValue().getName());
+                String extensionArrName = contextEntry.getValue().getName() + "Extensions";
+                childExtensionMap.putIfAbsent(extensionArrName, new HashSet<>());
+                DataTypesRegistry.getInstance().addDataType(extensionArrName);
+            }
+        }
+
+        for (Map.Entry<String, DatatypeTemplateContext> contextEntry : extensionDatatypeMap.entrySet()) {
+            /// Add Slice Extensions
+            /// Ensure that http:// keys are processed first
+            for (String key : searchKeys) {
+                if (contextEntry.getKey().contains(key.toLowerCase())) {
+                    childExtensionMap.get(key + "Extensions").add(contextEntry.getValue().getName());
+                }
+            }
+
+        }
+
+        childExtensionMap.entrySet().removeIf(entry -> entry.getValue().isEmpty());
+        extensionTemplateContext.setExtensionSlices(childExtensionMap);
+    }
+}

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/modelgen/AbstractPackageContextGenerator.java

@@ -79,6 +79,7 @@ private void populatePackageContext(Map<String, FHIRImplementationGuide> igEntri
 
             populateDatatypeTemplateContext(specificationData);
             populateResourceTemplateContext(implementationGuide);
+            populateExtensionTemplateContext(specificationData);
             populateIGTemplateContexts(entry.getValue().getName(), implementationGuide);
         }
         LOG.debug("Ended: Package Context population");
@@ -94,6 +95,8 @@ public BallerinaPackageGenToolConfig getToolConfig() {
 
     protected abstract void populateDatatypeTemplateContext(FHIRSpecificationData specificationData);
 
+    protected abstract void populateExtensionTemplateContext(FHIRSpecificationData specificationData);
+
     protected abstract void populateResourceTemplateContext(FHIRImplementationGuide implementationGuide);
 
     protected abstract void populateIGTemplateContexts(String igCode, FHIRImplementationGuide implementationGuide);

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/modelgen/AbstractResourceContextGenerator.java

@@ -111,17 +111,8 @@ protected void checkAndAddConstraintImport(Element element) {
      * @param element The element to check for international imports.
      */
     protected void checkAndAddInternationalImport(Element element) {
-        boolean isImportedFromInternational = false;
-        if (element.getContentReference() != null && GeneratorUtils.isReferredFromInternational(element.getContentReference())) {
-            isImportedFromInternational = true;
-        } else if (element.hasChildElements()) {
-            for (Element childElement : element.getChildElements().values()) {
-                if (childElement.getContentReference() != null && GeneratorUtils.isReferredFromInternational(childElement.getContentReference())) {
-                    isImportedFromInternational = true;
-                    break;
-                }
-            }
-        }
+        boolean isImportedFromInternational = GeneratorUtils.checkForInternationalImports(element);
+
         boolean isInternationalImportExists = this.resourceTemplateContextInstance.getResourceDependencies()
                 .stream()
                 .anyMatch(d -> d.equals(this.toolConfig.getPackageConfig().getInternationalPackage()));
@@ -183,7 +174,7 @@ private void validateAndPopulateExtendedElement(Element element) {
             extendedElement = GeneratorUtils.getInstance().populateExtendedElement(element, BallerinaDataType.Enum, elementDataType,
                     this.resourceTemplateContextInstance.getResourceName());
             putExtendedElementIfAbsent(element, extendedElement);
-        } else if (element.isSlice() || elementDataType.equals("BackboneElement") || elementDataType.equals("BackboneType") || (element.isExtended() && element.hasChildElements())) {
+        } else if (element.isSlice() || elementDataType.equals("BackboneElement") || elementDataType.equals("BackboneType") || (element.isExtended() && element.hasChildElements() && !GeneratorUtils.isPrimitiveElement(elementDataType))) {
             extendedElement = GeneratorUtils.getInstance().populateExtendedElement(element, BallerinaDataType.Record, elementDataType,
                     this.resourceTemplateContextInstance.getResourceName());
             extendedElement.setElements(element.getChildElements());

native/fhir-to-bal-lib/src/main/java/org/wso2/healthcare/fhir/ballerina/packagegen/tool/modelgen/versions/r4/R4DatatypeContextGenerator.java

@@ -63,7 +63,12 @@ protected void populateDatatypeContext() {
             } else if ("draft".equalsIgnoreCase(datatypeDefn.getDefinition().getStatus().toString()) &&
                     (datatypeDefn.getDefinition().getName().contains(" ") &&
                             datatypeDefn.getDefinition().getName().split(" ").length > 0)) {
-                continue;
+                String[] elementNameContent = datatypeDefn.getDefinition().getName().split(" ");
+
+                // Excludes generation of DataElement constraints (e.g.: ElementDefinition-de)
+                if (Arrays.asList(elementNameContent).contains("DataElement")) {
+                    continue;
+                }
             }
 
             DatatypeTemplateContext context = new DatatypeTemplateContext();
@@ -119,7 +124,7 @@ protected void populateDatatypeContext() {
                     context.addElement(element);
                 }
             }
-            datatypeTemplateContextMap().putIfAbsent(datatypeDefn.getDefinition().getUrl(), context);
+            getDatatypeTemplateContextMap().putIfAbsent(datatypeDefn.getDefinition().getUrl(), context);
         }
     }
 }