From 657c8045315d869654e563dc3b5b80598f8ffb1b Mon Sep 17 00:00:00 2001
From: David Pokora <dpokora@gmail.com>
Date: Mon, 14 Jan 2019 18:07:56 -0500
Subject: [PATCH] Immediate function inheritance support for compact-ast (TODO:
 legacy AST)

---
 slither/core/declarations/contract.py         | 11 ++++++++++-
 slither/solc_parsing/declarations/contract.py | 12 ++++++++++++
 slither/solc_parsing/slitherSolc.py           | 11 ++++++++++-
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/slither/core/declarations/contract.py b/slither/core/declarations/contract.py
index 02659b643..3bc56c938 100644
--- a/slither/core/declarations/contract.py
+++ b/slither/core/declarations/contract.py
@@ -20,6 +20,7 @@ class Contract(ChildSlither, SourceMapping):
         self._name = None
         self._id = None
         self._inheritance = []
+        self._immediate_inheritance = []
 
         self._enums = {}
         self._structures = {}
@@ -60,6 +61,13 @@ class Contract(ChildSlither, SourceMapping):
         '''
         return list(self._inheritance)
 
+    @property
+    def immediate_inheritance(self):
+        '''
+            list(Contract): List of contracts immediately inherited from (fathers). Order: order of declaration.
+        '''
+        return list(self._immediate_inheritance)
+
     @property
     def inheritance_reverse(self):
         '''
@@ -67,8 +75,9 @@ class Contract(ChildSlither, SourceMapping):
         '''
         return reversed(self._inheritance)
 
-    def setInheritance(self, inheritance):
+    def setInheritance(self, inheritance, immediate_inheritance):
         self._inheritance = inheritance
+        self._immediate_inheritance = immediate_inheritance
 
     @property
     def derived_contracts(self):
diff --git a/slither/solc_parsing/declarations/contract.py b/slither/solc_parsing/declarations/contract.py
index 29d951c1a..c0ac08db5 100644
--- a/slither/solc_parsing/declarations/contract.py
+++ b/slither/solc_parsing/declarations/contract.py
@@ -91,6 +91,18 @@ class ContractSolc04(Contract):
         self.linearizedBaseContracts = attributes['linearizedBaseContracts']
         self.fullyImplemented = attributes['fullyImplemented']
 
+        # Parse base contracts (immediate, non-linearized)
+        self.baseContracts = []
+        if self.is_compact_ast:
+            if 'baseContracts' in attributes:
+                for base_contract in attributes['baseContracts']:
+                    if base_contract['nodeType'] == 'InheritanceSpecifier':
+                        if 'baseName' in base_contract and 'referencedDeclaration' in base_contract['baseName']:
+                            self.baseContracts.append(base_contract['baseName']['referencedDeclaration'])
+        else:
+            # TODO: Parse from legacy-ast. 'baseContracts' is unreliable here. Possibly use 'children'.
+            pass
+
         # trufle does some re-mapping of id
         if 'baseContracts' in self._data:
             for elem in self._data['baseContracts']:
diff --git a/slither/solc_parsing/slitherSolc.py b/slither/solc_parsing/slitherSolc.py
index 39786d2db..9b8b7ad58 100644
--- a/slither/solc_parsing/slitherSolc.py
+++ b/slither/solc_parsing/slitherSolc.py
@@ -159,9 +159,18 @@ class SlitherSolc(Slither):
         # Update of the inheritance 
         for contract in self._contractsNotParsed:
             # remove the first elem in linearizedBaseContracts as it is the contract itself
+            ancestors = []
             fathers = []
             try:
+                # Resolve linearized base contracts.
                 for i in contract.linearizedBaseContracts[1:]:
+                    if i in contract.remapping:
+                        ancestors.append(self.get_contract_from_name(contract.remapping[i]))
+                    else:
+                        ancestors.append(self._contracts_by_id[i])
+
+                # Resolve immediate base contracts
+                for i in contract.baseContracts:
                     if i in contract.remapping:
                         fathers.append(self.get_contract_from_name(contract.remapping[i]))
                     else:
@@ -172,7 +181,7 @@ class SlitherSolc(Slither):
                 logger.error(red('Please read https://github.com/trailofbits/slither/wiki#keyerror-or-nonetype-error'))
                 logger.error(red('And update your code to remove the duplicate'))
                 exit(-1)
-            contract.setInheritance(fathers)
+            contract.setInheritance(ancestors, fathers)
 
         contracts_to_be_analyzed = self.contracts