From a7ebec271901a191edc24d22b991f29d66ea9a52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Per=20=C3=96stlund?= <perost86@gmail.com>
Date: Wed, 20 Nov 2019 15:03:21 +0100
Subject: [PATCH] [NF] Improve Inst.isDiscreteClass.

- Handle classes extended from builtin types better.
---
 OMCompiler/Compiler/NFFrontEnd/NFInst.mo      | 28 ++++++++++++---
 .../flattening/modelica/scodeinst/Makefile    |  1 +
 .../modelica/scodeinst/TypeExtends2.mo        | 34 +++++++++++++++++++
 3 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 testsuite/flattening/modelica/scodeinst/TypeExtends2.mo

diff --git a/OMCompiler/Compiler/NFFrontEnd/NFInst.mo b/OMCompiler/Compiler/NFFrontEnd/NFInst.mo
index 667b8622820..c6a8873e3bc 100644
--- a/OMCompiler/Compiler/NFFrontEnd/NFInst.mo
+++ b/OMCompiler/Compiler/NFFrontEnd/NFInst.mo
@@ -1951,15 +1951,35 @@ function updateComponentVariability
 protected
   Variability var = attr.variability;
 algorithm
-  if referenceEq(attr, NFComponent.DEFAULT_ATTR) and
-     Type.isDiscrete(Class.getType(cls, clsNode)) then
+  if referenceEq(attr, NFComponent.DEFAULT_ATTR) and isDiscreteClass(clsNode) then
     attr := NFComponent.IMPL_DISCRETE_ATTR;
-  elseif var == Variability.CONTINUOUS and
-     Type.isDiscrete(Class.getType(cls, clsNode)) then
+  elseif var == Variability.CONTINUOUS and isDiscreteClass(clsNode) then
     attr.variability := Variability.IMPLICITLY_DISCRETE;
   end if;
 end updateComponentVariability;
 
+function isDiscreteClass
+  input InstNode clsNode;
+  output Boolean isDiscrete;
+protected
+  InstNode base_node;
+  Class cls;
+  array<InstNode> exts;
+algorithm
+  base_node := Class.lastBaseClass(clsNode);
+  cls := InstNode.getClass(base_node);
+
+  isDiscrete := match cls
+    case Class.EXPANDED_CLASS(restriction = Restriction.TYPE())
+      algorithm
+        exts := ClassTree.getExtends(cls.elements);
+      then
+        if arrayLength(exts) == 1 then isDiscreteClass(exts[1]) else false;
+
+    else Type.isDiscrete(Class.getType(cls, base_node));
+  end match;
+end isDiscreteClass;
+
 function instTypeSpec
   input Absyn.TypeSpec typeSpec;
   input Modifier modifier;
diff --git a/testsuite/flattening/modelica/scodeinst/Makefile b/testsuite/flattening/modelica/scodeinst/Makefile
index f6fe01dd397..f595a78bec3 100644
--- a/testsuite/flattening/modelica/scodeinst/Makefile
+++ b/testsuite/flattening/modelica/scodeinst/Makefile
@@ -826,6 +826,7 @@ TypeDim3.mo \
 TypeDim4.mo \
 TypeDimNonType1.mo \
 TypeExtends1.mo \
+TypeExtends2.mo \
 TypeMissingBaseType1.mo \
 TypenameInvalid1.mo \
 TypenameInvalid2.mo \
diff --git a/testsuite/flattening/modelica/scodeinst/TypeExtends2.mo b/testsuite/flattening/modelica/scodeinst/TypeExtends2.mo
new file mode 100644
index 00000000000..099d8bdda5d
--- /dev/null
+++ b/testsuite/flattening/modelica/scodeinst/TypeExtends2.mo
@@ -0,0 +1,34 @@
+// name: TypeExtends2
+// keywords:
+// status: correct
+// cflags: -d=newInst
+//
+
+type TypeInteger
+  extends Integer;
+end TypeInteger;
+
+type TypeInteger2
+  extends TypeInteger;
+end TypeInteger2;
+
+type Color = TypeInteger2[3](each min = 0, each max = 255);
+
+model TypeExtends2
+  Color color = {0, 0, 0};
+  Integer icolor[3] = color;
+end TypeExtends2;
+
+// Result:
+// class TypeExtends2
+//   Integer color[1](min = 0, max = 255);
+//   Integer color[2](min = 0, max = 255);
+//   Integer color[3](min = 0, max = 255);
+//   Integer icolor[1];
+//   Integer icolor[2];
+//   Integer icolor[3];
+// equation
+//   color = {0, 0, 0};
+//   icolor = color;
+// end TypeExtends2;
+// endResult