From 1add42b691084c45125602dd38b582d3236f4d4e Mon Sep 17 00:00:00 2001
From: Brown <github@muglug.com>
Date: Mon, 25 Nov 2019 11:27:53 -0500
Subject: [PATCH] Fix switch handling

---
 .../plugins/PreventFloatAssignmentChecker.php |  2 +-
 .../Statements/Block/SwitchAnalyzer.php       | 14 ++++++--
 .../Internal/Visitor/TypeMappingVisitor.php   | 32 +++++++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)
 create mode 100644 src/Psalm/Internal/Visitor/TypeMappingVisitor.php

diff --git a/examples/plugins/PreventFloatAssignmentChecker.php b/examples/plugins/PreventFloatAssignmentChecker.php
index b3a292e0c..434cedfc4 100644
--- a/examples/plugins/PreventFloatAssignmentChecker.php
+++ b/examples/plugins/PreventFloatAssignmentChecker.php
@@ -34,7 +34,7 @@ class PreventFloatAssignmentChecker implements AfterExpressionAnalysisInterface
         array &$file_replacements = []
     ) {
         if ($expr instanceof PhpParser\Node\Expr\Assign
-            && ($expr_type = $statements_source->getNodeTypeProvider()->getType($expr))
+            && ($expr_type = $statements_source->getNodeTypeProvider()->getType($expr->expr))
             && $expr_type->hasFloat()
         ) {
             if (\Psalm\IssueBuffer::accepts(
diff --git a/src/Psalm/Internal/Analyzer/Statements/Block/SwitchAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Block/SwitchAnalyzer.php
index 9dc5e5cc6..e63f640ed 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Block/SwitchAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Block/SwitchAnalyzer.php
@@ -444,8 +444,6 @@ class SwitchAnalyzer
             );
         }
 
-        $statements_analyzer->node_data = $old_node_data;
-
         if ($switch_scope->negated_clauses) {
             $entry_clauses = Algebra::simplifyCNF(
                 array_merge(
@@ -537,6 +535,18 @@ class SwitchAnalyzer
 
         $statements_analyzer->analyze($case_stmts, $case_context);
 
+        $traverser = new PhpParser\NodeTraverser;
+        $traverser->addVisitor(
+            new \Psalm\Internal\Visitor\TypeMappingVisitor(
+                $statements_analyzer->node_data,
+                $old_node_data
+            )
+        );
+
+        $traverser->traverse([$case]);
+
+        $statements_analyzer->node_data = $old_node_data;
+
         /** @var array<string, bool> */
         $new_case_assigned_var_ids = $case_context->assigned_var_ids;
         $case_context->assigned_var_ids = $pre_assigned_var_ids + $new_case_assigned_var_ids;
diff --git a/src/Psalm/Internal/Visitor/TypeMappingVisitor.php b/src/Psalm/Internal/Visitor/TypeMappingVisitor.php
new file mode 100644
index 000000000..76026f74f
--- /dev/null
+++ b/src/Psalm/Internal/Visitor/TypeMappingVisitor.php
@@ -0,0 +1,32 @@
+<?php
+declare(strict_types=1);
+namespace Psalm\Internal\Visitor;
+
+use function array_map;
+use PhpParser\Node;
+use PhpParser\NodeVisitorAbstract;
+
+class TypeMappingVisitor extends NodeVisitorAbstract
+{
+    private $fake_type_provider;
+    private $real_type_provider;
+
+    public function __construct(
+        \Psalm\Internal\Provider\NodeDataProvider $fake_type_provider,
+        \Psalm\Internal\Provider\NodeDataProvider $real_type_provider
+    ) {
+        $this->fake_type_provider = $fake_type_provider;
+        $this->real_type_provider = $real_type_provider;
+    }
+
+    public function enterNode(Node $origNode)
+    {
+        /** @psalm-suppress ArgumentTypeCoercion */
+        $node_type = $this->fake_type_provider->getType($origNode);
+
+        if ($node_type) {
+            /** @psalm-suppress ArgumentTypeCoercion */
+            $this->real_type_provider->setType($origNode, clone $node_type);
+        }
+    }
+}