From b7432b80f7097f902ec35824e398f1c1e419e37b Mon Sep 17 00:00:00 2001
From: Tommy <info@flexforge.me>
Date: Thu, 21 Apr 2022 17:12:31 +0200
Subject: [PATCH] Improve InlineFunctionSqlCodeAnalysisRule

---
 .../InlineFunctionSqlCodeAnalysisRule.cs      | 48 ++++++++++++++-----
 .../dbx_codeanalysis_error_037.sql            |  1 +
 .../InlineFunctionSqlCodeAnalysisRule.xml     |  4 +-
 3 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/src/Dibix.Sdk/CodeAnalysis/Rules/InlineFunctionSqlCodeAnalysisRule.cs b/src/Dibix.Sdk/CodeAnalysis/Rules/InlineFunctionSqlCodeAnalysisRule.cs
index 05933841..ddb7a27a 100644
--- a/src/Dibix.Sdk/CodeAnalysis/Rules/InlineFunctionSqlCodeAnalysisRule.cs
+++ b/src/Dibix.Sdk/CodeAnalysis/Rules/InlineFunctionSqlCodeAnalysisRule.cs
@@ -40,6 +40,7 @@ public override void Visit(CreateFunctionStatement node)
 
         public override void Visit(CheckConstraintDefinition node)
         {
+            // CONSTRAINT [CK_dbx_codeanalysis_error_037_table_x] CHECK ([dbo].[dbx_codeanalysis_error_037_scalar]() = 0)
             ScalarFunctionCallVisitor visitor = new ScalarFunctionCallVisitor(this.IsScalarFunctionCall);
             node.CheckCondition.Accept(visitor);
             foreach (FunctionCall call in visitor.Locations)
@@ -51,6 +52,14 @@ public override void Visit(CheckConstraintDefinition node)
 
         public override void Visit(DeclareVariableElement node)
         {
+            // DECLARE @y INT = [dbo].[dbx_codeanalysis_error_037_scalar]()
+            if (node.Value is FunctionCall call)
+            {
+                VisitFunctionCall(call);
+                return;
+            }
+
+            // DECLARE @x INT = (SELECT [dbo].[dbx_codeanalysis_error_037_scalar]())
             if (node.Value is ScalarSubquery scalarSubquery
              && scalarSubquery.QueryExpression is QuerySpecification querySpecification)
                 this.VisitQuerySpecification(querySpecification);
@@ -58,32 +67,47 @@ public override void Visit(DeclareVariableElement node)
 
         public override void Visit(SelectStatement node)
         {
+            // SELECT @x = [dbo].[dbx_codeanalysis_error_037_scalar]()
+            // and
+            // SELECT [dbo].[dbx_codeanalysis_error_037_scalar]()
             if (node.QueryExpression is QuerySpecification querySpecification)
                 this.VisitQuerySpecification(querySpecification);
         }
 
         public override void Visit(SetVariableStatement node)
         {
-            TSqlFragment target = node?.Expression;
-            if (target == null)
-                return;
-
-            if (this._scalarFunctionCalls.ContainsKey(target.StartOffset))
-                this._scalarFunctionCalls.Remove(target.StartOffset);
+            // SET @x = [dbo].[dbx_codeanalysis_error_037_scalar]()
+            if (node?.Expression is FunctionCall call)
+                VisitFunctionCall(call);
         }
 
         private void VisitQuerySpecification(QuerySpecification node)
         {
-            if (node.SelectElements.Count == 1
-             && node.FromClause == null)
+            if (node.SelectElements.Count != 1 || node.FromClause != null) 
+                return;
+
+            TSqlFragment target = node.SelectElements[0];
+            switch (target)
             {
-                TSqlFragment target = node.SelectElements[0];
-                if (target is SelectSetVariable selectSetVariable)
+                // SELECT @x = [dbo].[dbx_codeanalysis_error_037_scalar]()
+                case SelectSetVariable selectSetVariable:
                     target = selectSetVariable.Expression;
+                    break;
 
-                if (this._scalarFunctionCalls.ContainsKey(target.StartOffset))
-                    this._scalarFunctionCalls.Remove(target.StartOffset);
+                // DECLARE @x INT = (SELECT [dbo].[dbx_codeanalysis_error_037_scalar]())
+                case SelectScalarExpression selectScalarExpression:
+                    target = selectScalarExpression.Expression;
+                    break;
             }
+
+            if (target is FunctionCall call)
+                this.VisitFunctionCall(call);
+        }
+
+        private void VisitFunctionCall(FunctionCall call)
+        {
+            if (this._scalarFunctionCalls.ContainsKey(call.StartOffset) && IsScalarFunctionCall(call))
+                this._scalarFunctionCalls.Remove(call.StartOffset);
         }
 
         private bool IsScalarFunctionCall(FunctionCall call) => base.Model.IsScalarFunction(call);
diff --git a/tests/Dibix.Sdk.Tests.Database/CodeAnalysis/dbx_codeanalysis_error_037.sql b/tests/Dibix.Sdk.Tests.Database/CodeAnalysis/dbx_codeanalysis_error_037.sql
index 3348c956..4bdd4065 100644
--- a/tests/Dibix.Sdk.Tests.Database/CodeAnalysis/dbx_codeanalysis_error_037.sql
+++ b/tests/Dibix.Sdk.Tests.Database/CodeAnalysis/dbx_codeanalysis_error_037.sql
@@ -31,6 +31,7 @@ AS
 BEGIN
 	
 	DECLARE @x INT = (SELECT [dbo].[dbx_codeanalysis_error_037_scalar]()) -- OK
+	DECLARE @y INT = [dbo].[dbx_codeanalysis_error_037_scalar]() -- OK
 	SELECT @x = [dbo].[dbx_codeanalysis_error_037_scalar]() -- OK
 	SET @x = [dbo].[dbx_codeanalysis_error_037_scalar]() -- OK
 	SELECT [dbo].[dbx_codeanalysis_error_037_scalar]() -- OK
diff --git a/tests/Dibix.Sdk.Tests/Resources/CodeAnalysis/InlineFunctionSqlCodeAnalysisRule.xml b/tests/Dibix.Sdk.Tests/Resources/CodeAnalysis/InlineFunctionSqlCodeAnalysisRule.xml
index c002ebbb..085652c0 100644
--- a/tests/Dibix.Sdk.Tests/Resources/CodeAnalysis/InlineFunctionSqlCodeAnalysisRule.xml
+++ b/tests/Dibix.Sdk.Tests/Resources/CodeAnalysis/InlineFunctionSqlCodeAnalysisRule.xml
@@ -8,11 +8,11 @@
   <error
     ruleid="37"
     message="Scalar functions should only be used for assignments or check constraints. Otherwise replace it with an inline table-valued function: dbx_codeanalysis_error_037_scalar"
-    line="40"
+    line="41"
     column="16" />
   <error
     ruleid="37"
     message="Scalar functions should only be used for assignments or check constraints. Otherwise replace it with an inline table-valued function: dbx_codeanalysis_error_037_scalar"
-    line="42"
+    line="43"
     column="18" />
 </errors>
\ No newline at end of file